diff --git a/Cargo.toml b/Cargo.toml index 42aa113c1..cd872a17f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ async-trait = "0.1" connection-string = "0.2" num-traits = "0.2" uuid = "1.0" +zeroize = "1.8.2" [target.'cfg(windows)'.dependencies] winauth = { version = "0.0.4", optional = true } diff --git a/docker/test-server.sh b/docker/test-server.sh new file mode 100755 index 000000000..a0239280c --- /dev/null +++ b/docker/test-server.sh @@ -0,0 +1,86 @@ +#!/usr/bin/env bash +# +# Start a SQL Server for the test suite, with podman or docker. +# +# ./docker/test-server.sh up # build, start, wait until it accepts connections +# ./docker/test-server.sh down # stop and remove +# ./docker/test-server.sh logs # follow the server log +# +# Then: +# +# export TIBERIUS_TEST_CONNECTION_STRING='server=tcp:localhost,1433;user=SA;password=;IntegratedSecurity=true;TrustServerCertificate=true' +# cargo test +# +# IMAGE selects the flavour; the default works on both x86_64 and arm64. +# The full SQL Server images are x86_64 only, so on an arm64 machine +# (Apple silicon) they either refuse to run or run under emulation. + +set -euo pipefail + +ENGINE="${ENGINE:-$(command -v podman >/dev/null 2>&1 && echo podman || echo docker)}" +NAME="${NAME:-tiberius-test-mssql}" +PORT="${PORT:-1433}" +IMAGE="${IMAGE:-azure-sql-edge}" +PASSWORD='' +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +case "${1:-up}" in + up) + echo "engine: $ENGINE image: $IMAGE port: $PORT" + "$ENGINE" build -q -f "$HERE/docker-$IMAGE.dockerfile" -t "$NAME:local" "$HERE" + "$ENGINE" rm -f "$NAME" >/dev/null 2>&1 || true + "$ENGINE" run -d --name "$NAME" \ + -e ACCEPT_EULA=Y \ + -e "MSSQL_SA_PASSWORD=$PASSWORD" \ + -e "SA_PASSWORD=$PASSWORD" \ + -p "$PORT:1433" \ + "$NAME:local" >/dev/null + + # The port opens well before the server will answer, so poll the log + # rather than the socket. + # + # The log is captured into a variable and matched there, rather than + # piped into `grep -q`. Under `set -o pipefail`, `grep -q` exits on the + # first match, the writer upstream dies of SIGPIPE, and the pipeline + # reports failure even though the match succeeded — so the wait never + # ends. + echo -n "waiting for SQL Server" + for _ in $(seq 1 120); do + logs="$("$ENGINE" logs "$NAME" 2>&1 || true)" + + case "$logs" in + *"SQL Server is now ready for client connections"*) + echo " — ready" + exit 0 + ;; + esac + + running="$("$ENGINE" ps --format '{{.Names}}' || true)" + case "$running" in + *"$NAME"*) ;; + *) + echo " — container exited:" + "$ENGINE" logs --tail 30 "$NAME" || true + exit 1 + ;; + esac + + echo -n . + sleep 2 + done + echo " — gave up; last lines:" + "$ENGINE" logs --tail 30 "$NAME" + exit 1 + ;; + down) + "$ENGINE" rm -f "$NAME" >/dev/null 2>&1 || true + echo "removed $NAME" + ;; + logs) + "$ENGINE" logs -f "$NAME" + ;; + *) + echo "usage: $0 {up|down|logs}" >&2 + exit 2 + ;; +esac diff --git a/src/client.rs b/src/client.rs index 688721d10..fd30a2dd0 100644 --- a/src/client.rs +++ b/src/client.rs @@ -251,10 +251,13 @@ impl Client { Ok(result) } - /// Execute a `BULK INSERT` statement, efficiantly storing a large number of + /// Execute a `BULK INSERT` statement, efficiently storing a large number of /// rows to a specified table. Note: make sure the input row follows the same /// schema as the table, otherwise calling `send()` will return an error. /// + /// This is equivalent to calling `bulk_insert("table_name", &["*"])` to merge + /// all of a tables columns. + /// /// # Example /// /// ``` @@ -299,12 +302,67 @@ impl Client { pub async fn bulk_insert<'a>( &'a mut self, table: &'a str, + ) -> crate::Result> { + self.bulk_insert_columns(table, &["*"]).await + } + + /// Execute a `BULK INSERT` statement, efficiently storing a large number of + /// rows to a specified table. Note: make sure the input row follows the same + /// schema as the column list, otherwise calling `send()` will return an error. + /// + /// # Example + /// + /// ``` + /// # use tiberius::{Config, IntoRow}; + /// # 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 create_table = r#" + /// CREATE TABLE ##bulk_test ( + /// id INT IDENTITY PRIMARY KEY, + /// foo INT NOT NULL, + /// bar FLOAT NOT NULL + /// ) + /// "#; + /// + /// client.simple_query(create_table).await?; + /// + /// // Start the bulk insert with the client. + /// let mut req = client.bulk_insert_columns("##bulk_test", &["foo", "bar"]).await?; + /// + /// for (i, j) in [(0i32, 0f64), (1i32, 1f64), (2i32, 2f64)] { + /// let row = (i, j).into_row(); + /// + /// // The request will handle flushing to the wire in an optimal way, + /// // balancing between memory usage and IO performance. + /// req.send(row).await?; + /// } + /// + /// // The request must be finalized. + /// let res = req.finalize().await?; + /// assert_eq!(3, res.total()); + /// # Ok(()) + /// # } + /// ``` + pub async fn bulk_insert_columns<'a>( + &'a mut self, + table: &'a str, + columns: &'a [&'a str], ) -> crate::Result> { // Start the bulk request self.connection.flush_stream().await?; // retrieve column metadata from server - let query = format!("SELECT TOP 0 * FROM {}", table); + let columns = columns.join(", "); + let query = format!("SELECT TOP 0 {columns} FROM {table}"); let req = BatchRequest::new(query, self.connection.context().transaction_descriptor()); diff --git a/src/client/auth.rs b/src/client/auth.rs index 208d8d060..3abf42df8 100644 --- a/src/client/auth.rs +++ b/src/client/auth.rs @@ -1,18 +1,15 @@ use std::fmt::Debug; +use zeroize::Zeroizing; #[derive(Clone, PartialEq, Eq)] pub struct SqlServerAuth { user: String, - password: String, + password: Zeroizing, } impl SqlServerAuth { - pub(crate) fn user(&self) -> &str { - &self.user - } - - pub(crate) fn password(&self) -> &str { - &self.password + pub(crate) fn into_credentials(self) -> (String, Zeroizing) { + (self.user, self.password) } } @@ -79,7 +76,7 @@ impl AuthMethod { pub fn sql_server(user: impl ToString, password: impl ToString) -> Self { Self::SqlServer(SqlServerAuth { user: user.to_string(), - password: password.to_string(), + password: Zeroizing::new(password.to_string()), }) } @@ -104,3 +101,25 @@ impl AuthMethod { Self::AADToken(token.to_string()) } } + +#[cfg(test)] +mod tests { + use super::AuthMethod; + use zeroize::Zeroize; + + #[test] + fn sql_server_password_can_be_consumed_and_zeroized() { + let AuthMethod::SqlServer(auth) = AuthMethod::sql_server("sa", "secret") else { + unreachable!(); + }; + + let (user, mut password) = auth.into_credentials(); + + assert_eq!("sa", user); + assert_eq!("secret", password.as_str()); + + password.zeroize(); + + assert!(password.is_empty()); + } +} diff --git a/src/client/config.rs b/src/client/config.rs index 57374f8ce..3d6994f0f 100644 --- a/src/client/config.rs +++ b/src/client/config.rs @@ -32,6 +32,7 @@ pub struct Config { pub(crate) trust: TrustConfig, pub(crate) auth: AuthMethod, pub(crate) readonly: bool, + pub(crate) packet_size: Option, } #[derive(Clone, Debug)] @@ -65,6 +66,7 @@ impl Default for Config { trust: TrustConfig::Default, auth: AuthMethod::None, readonly: false, + packet_size: None, } } } @@ -115,6 +117,22 @@ impl Config { self.application_name = Some(name.to_string()); } + /// Sets the TDS packet size for the connection. + /// + /// Larger packet sizes can improve bulk insert performance by reducing + /// the number of network round-trips. Valid values are 512 to 32767. + /// The server may negotiate a different size. + /// + /// - Defaults to 4096 bytes. + pub fn packet_size(&mut self, size: u32) { + self.packet_size = Some(size); + } + + /// Gets the configured packet size, if set. + pub fn get_packet_size(&self) -> Option { + self.packet_size + } + /// Set the preferred encryption level. /// /// - With `tls` feature, defaults to `Required`. diff --git a/src/client/connection.rs b/src/client/connection.rs index 0038386f4..ae136fcc8 100644 --- a/src/client/connection.rs +++ b/src/client/connection.rs @@ -20,7 +20,7 @@ use asynchronous_codec::Framed; use bytes::BytesMut; #[cfg(any(windows, feature = "integrated-auth-gssapi"))] use codec::TokenSspi; -use futures_util::io::{AsyncRead, AsyncWrite}; +use futures_util::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use futures_util::ready; use futures_util::sink::SinkExt; use futures_util::stream::{Stream, TryStream, TryStreamExt}; @@ -39,6 +39,7 @@ use task::Poll; use tracing::{event, Level}; #[cfg(all(windows, feature = "winauth"))] use winauth::{windows::NtlmSspiBuilder, NextBytes}; +use zeroize::{Zeroize, Zeroizing}; /// A `Connection` is an abstraction between the [`Client`] and the server. It /// can be used as a `Stream` to fetch [`Packet`]s from and to `send` packets @@ -106,6 +107,7 @@ impl Connection { config.host, config.application_name, config.readonly, + config.packet_size, prelogin, ) .await?; @@ -196,6 +198,45 @@ impl Connection { Ok(()) } + async fn send_sensitive_login( + &mut self, + mut header: PacketHeader, + mut payload: Zeroizing>, + ) -> crate::Result<()> { + self.flushed = false; + let packet_size = (self.context.packet_size() as usize) - HEADER_BYTES; + let mut offset = 0; + + while offset < payload.len() { + let end = cmp::min(payload.len(), offset + packet_size); + + if end == payload.len() { + header.set_status(PacketStatus::EndOfMessage); + } else { + header.set_status(PacketStatus::NormalMessage); + } + + let mut frame = Zeroizing::new(Vec::with_capacity(HEADER_BYTES + end - offset)); + header.encode(&mut *frame)?; + frame.extend_from_slice(&payload[offset..end]); + + let size = (frame.len() as u16).to_be_bytes(); + frame[2] = size[0]; + frame[3] = size[1]; + + event!(Level::TRACE, "Sending a packet ({} bytes)", frame.len(),); + + self.transport.write_all(frame.as_slice()).await?; + frame.zeroize(); + payload[offset..end].zeroize(); + offset = end; + } + + (&mut *self.transport).flush().await?; + + Ok(()) + } + /// Sends a packet of data to the database. /// /// # Warning @@ -293,6 +334,7 @@ impl Connection { server_name: Option, application_name: Option, readonly: bool, + packet_size: Option, prelogin: PreloginMessage, ) -> crate::Result { let mut login_message = LoginMessage::new(); @@ -311,6 +353,10 @@ impl Connection { login_message.readonly(readonly); + if let Some(size) = packet_size { + login_message.packet_size(size); + } + match auth { #[cfg(all(windows, feature = "winauth"))] AuthMethod::Integrated => { @@ -332,7 +378,7 @@ impl Connection { event!(Level::TRACE, sspi_response_len = sspi_response.len()); let id = self.context.next_packet_id(); - let header = PacketHeader::login(id); + let header = PacketHeader::sspi(id); let token = TokenSspi::new(sspi_response); self.send(header, token).await?; @@ -415,11 +461,16 @@ impl Connection { self = self.post_login_encryption(encryption); } AuthMethod::SqlServer(auth) => { - login_message.user_name(auth.user()); - login_message.password(auth.password()); + let (user, mut password) = auth.into_credentials(); + + login_message.user_name(user); + login_message.password(password.as_str()); + let payload = login_message.encode_to_vec()?; + password.zeroize(); let id = self.context.next_packet_id(); - self.send(PacketHeader::login(id), login_message).await?; + self.send_sensitive_login(PacketHeader::login(id), payload) + .await?; self = self.post_login_encryption(encryption); } AuthMethod::AADToken(token) => { diff --git a/src/error.rs b/src/error.rs index 98bf01b58..504ca8c18 100644 --- a/src/error.rs +++ b/src/error.rs @@ -8,8 +8,8 @@ use thiserror::Error; /// the lifecycle of this driver #[derive(Debug, Clone, Error, PartialEq, Eq)] pub enum Error { - #[error("An error occured during the attempt of performing I/O: {}", message)] - /// An error occured when performing I/O to the server. + #[error("An error occurred during the attempt of performing I/O: {}", message)] + /// An error occurred when performing I/O to the server. Io { /// A list specifying general categories of I/O error. kind: IoErrorKind, diff --git a/src/query.rs b/src/query.rs index 790052b4d..f31c54b1c 100644 --- a/src/query.rs +++ b/src/query.rs @@ -35,6 +35,119 @@ impl<'a> Query<'a> { self.params.push(param.into_sql()); } + /// Bind every item of an iterator, in order. + /// + /// Equivalent to calling [`bind`] once per item. Pairs with + /// [`placeholders`] to build an `IN` list, where the number of + /// parameters is only known at runtime. + /// + /// # Example + /// + /// ``` + /// # use tiberius::Query; + /// let ids = vec![1i32, 2, 3]; + /// + /// let sql = format!( + /// "SELECT name FROM users WHERE id IN ({})", + /// Query::placeholders(1, ids.len()), + /// ); + /// + /// let mut query = Query::new(sql); + /// query.bind_iter(ids); + /// + /// assert_eq!(query.param_count(), 3); + /// ``` + /// + /// [`bind`]: #method.bind + /// [`placeholders`]: #method.placeholders + pub fn bind_iter(&mut self, params: impl IntoIterator + 'a>) { + for param in params { + self.bind(param); + } + } + + /// How many parameters have been bound so far. + /// + /// Useful for checking against [`MAX_PARAMETERS`] before executing a + /// statement whose parameter count is decided at runtime. + /// + /// [`MAX_PARAMETERS`]: #associatedconstant.MAX_PARAMETERS + pub fn param_count(&self) -> usize { + self.params.len() + } + + /// The largest number of parameters SQL Server accepts in one statement. + /// + /// A statement carrying more is rejected by the server with + /// "The incoming request has too many parameters. The server supports a + /// maximum of 2100 parameters." — which arrives only after the whole + /// batch has been sent. + /// + /// This matters most for an `IN` list or a multi-row `INSERT`, where the + /// count comes from the length of a collection rather than from the SQL + /// text: the limit is reached by data volume, at run time, on a batch + /// that may be larger than any that was tested. Split such a batch into + /// chunks of at most `MAX_PARAMETERS / parameters_per_row` items. + /// + /// # Example + /// + /// ``` + /// # use tiberius::Query; + /// // A three-column INSERT: three parameters per row. + /// let rows_per_statement = Query::MAX_PARAMETERS / 3; + /// assert_eq!(rows_per_statement, 700); + /// ``` + pub const MAX_PARAMETERS: usize = 2100; + + /// Build a `@P1, @P2, …` placeholder list for `count` parameters, + /// numbered from `first`. + /// + /// SQL Server has no array parameter, so an `IN` list must name one + /// placeholder per value, and `IN (@P1)` bound to a comma-separated + /// string matches nothing rather than failing. Generating the list is + /// the only way to write such a query, and this does it without a + /// format loop at every call site. + /// + /// `first` is 1-based, matching the `@P1` numbering + /// [`Query::new`] documents. + /// + /// # Example + /// + /// ``` + /// # use tiberius::Query; + /// assert_eq!(Query::placeholders(1, 3), "@P1, @P2, @P3"); + /// + /// // Continuing after parameters that are already bound. + /// assert_eq!(Query::placeholders(4, 2), "@P4, @P5"); + /// ``` + /// + /// A count of zero yields an empty string. `IN ()` is a syntax error, so + /// a caller with nothing to match on should skip the query rather than + /// build one: + /// + /// ``` + /// # use tiberius::Query; + /// let ids: Vec = Vec::new(); + /// assert!(Query::placeholders(1, ids.len()).is_empty()); + /// ``` + /// + /// [`Query::new`]: #method.new + pub fn placeholders(first: usize, count: usize) -> String { + use std::fmt::Write; + + let mut out = String::with_capacity(count * 6); + + for index in 0..count { + if index > 0 { + out.push_str(", "); + } + // Writing into a String cannot fail. + let _ = write!(out, "@P{}", first + index); + } + + out + } + /// Executes SQL statements in the SQL Server, returning the number rows /// affected. Useful for `INSERT`, `UPDATE` and `DELETE` statements. See /// [`Client#execute`] for a simpler API if the parameters are statically @@ -136,3 +249,75 @@ impl<'a> Query<'a> { Ok(result) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn placeholders_are_numbered_from_one() { + assert_eq!(Query::placeholders(1, 1), "@P1"); + assert_eq!(Query::placeholders(1, 3), "@P1, @P2, @P3"); + } + + #[test] + fn placeholders_can_continue_from_an_offset() { + // For a query that already binds parameters before the list. + assert_eq!(Query::placeholders(4, 2), "@P4, @P5"); + assert_eq!(Query::placeholders(10, 1), "@P10"); + } + + #[test] + fn no_placeholders_is_an_empty_string() { + // `IN ()` is a syntax error, so a caller with nothing to match on + // must skip the query rather than build one. + assert_eq!(Query::placeholders(1, 0), ""); + assert_eq!(Query::placeholders(7, 0), ""); + } + + #[test] + fn placeholders_have_no_trailing_separator() { + let list = Query::placeholders(1, 5); + assert!(!list.ends_with(", ")); + assert_eq!(list.matches(',').count(), 4); + } + + #[test] + fn binding_an_iterator_counts_every_item() { + let mut query = Query::new("SELECT 1"); + assert_eq!(query.param_count(), 0); + + query.bind_iter(vec![1i32, 2, 3]); + assert_eq!(query.param_count(), 3); + + query.bind(4i32); + assert_eq!(query.param_count(), 4); + } + + #[test] + fn binding_an_empty_iterator_binds_nothing() { + let mut query = Query::new("SELECT 1"); + query.bind_iter(Vec::::new()); + assert_eq!(query.param_count(), 0); + } + + #[test] + fn a_generated_list_matches_the_number_of_bound_parameters() { + // The invariant that makes this pair usable: one placeholder per + // bound value, or the server rejects the statement. + let ids = vec![10i32, 20, 30, 40]; + let list = Query::placeholders(1, ids.len()); + + let mut query = Query::new(format!("SELECT * FROM t WHERE id IN ({list})")); + query.bind_iter(ids); + + assert_eq!(list.matches("@P").count(), query.param_count()); + } + + #[test] + fn the_parameter_limit_is_the_documented_tds_maximum() { + assert_eq!(Query::MAX_PARAMETERS, 2100); + // The chunking arithmetic the docs describe. + assert_eq!(Query::MAX_PARAMETERS / 3, 700); + } +} diff --git a/src/row.rs b/src/row.rs index 7eaf79ff2..3610b5818 100644 --- a/src/row.rs +++ b/src/row.rs @@ -411,14 +411,23 @@ impl Row { where R: FromSql<'a>, I: QueryIdx, + { + let data = self.get_column_data(idx)?; + + R::from_sql(data) + } + + /// Retrieve a column's data for a given column index. + #[track_caller] + pub fn get_column_data(&self, idx: I) -> crate::Result<&ColumnData<'static>> + where + I: QueryIdx, { let idx = idx.idx(self).ok_or_else(|| { Error::Conversion(format!("Could not find column with index {}", idx).into()) })?; - let data = self.data.get(idx).unwrap(); - - R::from_sql(data) + Ok(self.data.get(idx).unwrap()) } } diff --git a/src/tds/codec/header.rs b/src/tds/codec/header.rs index cfabf2c55..1ceda7556 100644 --- a/src/tds/codec/header.rs +++ b/src/tds/codec/header.rs @@ -92,6 +92,17 @@ impl PacketHeader { } } + // Only constructed on auth code paths gated behind platform/feature cfgs + // (Windows winauth / integrated-auth-gssapi), so it reads as dead on other builds. + #[allow(dead_code)] + pub fn sspi(id: u8) -> Self { + Self { + ty: PacketType::Sspi, + status: PacketStatus::EndOfMessage, + ..Self::new(0, id) + } + } + pub fn batch(id: u8) -> Self { Self { ty: PacketType::SQLBatch, diff --git a/src/tds/codec/login.rs b/src/tds/codec/login.rs index bab14d736..c92684f71 100644 --- a/src/tds/codec/login.rs +++ b/src/tds/codec/login.rs @@ -5,6 +5,7 @@ use enumflags2::{bitflags, BitFlags}; use io::{Cursor, Write}; use std::fmt::Debug; use std::{borrow::Cow, io}; +use zeroize::{Zeroize, Zeroizing}; uint_enum! { #[repr(u32)] @@ -235,10 +236,16 @@ impl<'a> LoginMessage<'a> { self.type_flags.remove(LoginTypeFlag::ReadOnlyIntent); } } -} -impl<'a> Encode for LoginMessage<'a> { - fn encode(self, dst: &mut BytesMut) -> crate::Result<()> { + /// Sets the requested TDS packet size. + /// + /// Valid values are 512 to 32767. The server may negotiate a different size. + /// Larger packet sizes can improve bulk insert performance. + pub fn packet_size(&mut self, size: u32) { + self.packet_size = size; + } + + pub(crate) fn encode_to_vec(self) -> crate::Result>> { let mut cursor = Cursor::new(Vec::with_capacity(512)); // Space for the length @@ -361,7 +368,7 @@ impl<'a> Encode for LoginMessage<'a> { for codepoint in fed_auth_ext.fed_auth_token.encode_utf16() { token.write_u16::(codepoint)?; } - let token = token.into_inner(); + let mut token = token.into_inner(); // options (1) + TokenLength(4) + Token.length + nonce.length let feature_ext_length = @@ -378,6 +385,7 @@ impl<'a> Encode for LoginMessage<'a> { cursor.write_u32::(token.len() as u32)?; cursor.write_all(token.as_slice())?; + token.zeroize(); if let Some(nonce) = fed_auth_ext.nonce { cursor.write_all(nonce.as_ref())?; @@ -389,7 +397,15 @@ impl<'a> Encode for LoginMessage<'a> { cursor.set_position(0); cursor.write_u32::(cursor.get_ref().len() as u32)?; - dst.extend(cursor.into_inner()); + Ok(Zeroizing::new(cursor.into_inner())) + } +} + +impl<'a> Encode for LoginMessage<'a> { + fn encode(self, dst: &mut BytesMut) -> crate::Result<()> { + let mut encoded = self.encode_to_vec()?; + dst.extend_from_slice(encoded.as_slice()); + encoded.zeroize(); Ok(()) } diff --git a/src/tds/codec/token/token_col_metadata.rs b/src/tds/codec/token/token_col_metadata.rs index 53ffdf1c6..6d49938dd 100644 --- a/src/tds/codec/token/token_col_metadata.rs +++ b/src/tds/codec/token/token_col_metadata.rs @@ -25,7 +25,7 @@ pub struct MetaDataColumn<'a> { impl<'a> Display for MetaDataColumn<'a> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{} ", self.col_name)?; + write!(f, "[{}] ", self.col_name)?; match &self.base.ty { TypeInfo::FixedLen(fixed) => match fixed { diff --git a/src/tds/numeric.rs b/src/tds/numeric.rs index e4eff9ceb..f95e7cbad 100644 --- a/src/tds/numeric.rs +++ b/src/tds/numeric.rs @@ -186,9 +186,10 @@ impl Debug for Numeric { fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> { write!( f, - "{}.{:0pad$}", - self.int_part(), - self.dec_part(), + "{}{}.{:0pad$}", + if self.value() < 0 { "-" } else { "" }, + self.int_part().abs(), + self.dec_part().abs(), pad = self.scale as usize ) } @@ -263,6 +264,23 @@ mod decimal { Numeric::new_with_scale(value, self_.scale() as u8) }); ); + + #[cfg(feature = "tds73")] + into_sql!(self_, + Decimal: (ColumnData::Numeric, { + let unpacked = self_.unpack(); + + let mut value = (((unpacked.hi as u128) << 64) + + ((unpacked.mid as u128) << 32) + + unpacked.lo as u128) as i128; + + if self_.is_sign_negative() { + value = -value; + } + + Numeric::new_with_scale(value, self_.scale() as u8) + }); + ); } #[cfg(feature = "bigdecimal")] @@ -368,6 +386,36 @@ mod tests { assert_eq!(n.dec_part(), 5); } + #[test] + fn numeric_to_string() { + assert_eq!(Numeric::new_with_scale(123, 0).to_string(), "123.0"); + assert_eq!(Numeric::new_with_scale(123, 1).to_string(), "12.3"); + assert_eq!(Numeric::new_with_scale(123, 2).to_string(), "1.23"); + assert_eq!(Numeric::new_with_scale(123, 3).to_string(), "0.123"); + assert_eq!(Numeric::new_with_scale(123, 4).to_string(), "0.0123"); + assert_eq!( + Numeric::new_with_scale(123, 36).to_string(), + "0.000000000000000000000000000000000123" + ); + assert_eq!( + Numeric::new_with_scale(123, 37).to_string(), + "0.0000000000000000000000000000000000123" + ); + assert_eq!(Numeric::new_with_scale(-123, 0).to_string(), "-123.0"); + assert_eq!(Numeric::new_with_scale(-123, 1).to_string(), "-12.3"); + assert_eq!(Numeric::new_with_scale(-123, 2).to_string(), "-1.23"); + assert_eq!(Numeric::new_with_scale(-123, 3).to_string(), "-0.123"); + assert_eq!(Numeric::new_with_scale(-123, 4).to_string(), "-0.0123"); + assert_eq!( + Numeric::new_with_scale(-123, 36).to_string(), + "-0.000000000000000000000000000000000123" + ); + assert_eq!( + Numeric::new_with_scale(-123, 37).to_string(), + "-0.0000000000000000000000000000000000123" + ); + } + #[test] fn calculates_precision_correctly() { let n = Numeric::new_with_scale(57705, 2); diff --git a/src/tds/stream/query.rs b/src/tds/stream/query.rs index 0dc694749..8522cf4a0 100644 --- a/src/tds/stream/query.rs +++ b/src/tds/stream/query.rs @@ -222,28 +222,21 @@ impl<'a> QueryStream<'a> { /// of querying. pub async fn into_results(mut self) -> crate::Result>> { let mut results: Vec> = Vec::new(); - let mut result: Option> = None; + let mut result: Vec = if self.try_next().await?.is_some() { + Vec::new() + } else { + return Ok(results); + }; while let Some(item) = self.try_next().await? { - match (item, &mut result) { - (QueryItem::Row(row), None) => { - result = Some(vec![row]); - } - (QueryItem::Row(row), Some(ref mut result)) => result.push(row), - (QueryItem::Metadata(_), None) => { - result = Some(Vec::new()); - } - (QueryItem::Metadata(_), ref mut previous_result) => { - results.push(previous_result.take().unwrap()); - result = None; - } + if let QueryItem::Row(row) = item { + result.push(row); + } else { + results.push(result); + result = Vec::new(); } } - - if let Some(result) = result { - results.push(result); - } - + results.push(result); Ok(results) } diff --git a/tests/bulk.rs b/tests/bulk.rs index 33b90637a..bc1c2c8d3 100644 --- a/tests/bulk.rs +++ b/tests/bulk.rs @@ -4,6 +4,7 @@ use once_cell::sync::Lazy; use std::cell::RefCell; use std::env; use std::sync::Once; +use tiberius::ColumnData; use tiberius::{IntoSql, Result, TokenRow}; #[cfg(all(feature = "tds73", feature = "chrono"))] @@ -218,3 +219,134 @@ test_bulk_type!(datetime2_7( 100, vec![DateTime::from_timestamp(1658524194, 123456789); 100].into_iter() )); + +#[test_on_runtimes] +async fn read_and_write_to_keyword_columns(mut conn: tiberius::Client) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let table = format!("##{}", random_table().await); + + conn.simple_query(format!("CREATE TABLE {} ([End] INT)", table)) + .await?; + + let mut req = conn.bulk_insert(&table).await.unwrap(); + for num in [6, 7, 8] { + let mut row = TokenRow::new(); + row.push(ColumnData::I32(Some(num))); + req.send(row).await.unwrap(); + } + let result = req.finalize().await.unwrap(); + assert_eq!(result.rows_affected(), &[3]); + + let rows = conn + .query(format!("SELECT [End] FROM {}", table), &[]) + .await? + .into_first_result() + .await?; + + assert_eq!(rows.len(), 3); + assert_eq!(Some(6), rows[0].get(0)); + assert_eq!(Some(7), rows[1].get(0)); + assert_eq!(Some(8), rows[2].get(0)); + + Ok(()) +} + +macro_rules! test_bulk_columns { + ($name:ident($total_generated:literal $(, $sql_type:literal)+ $(, ($cols:expr, $generator:expr ))+ $(,)?)) => { + paste::item! { + #[test_on_runtimes] + async fn [< bulk_load_optional_ $name >](mut conn: tiberius::Client) -> Result<()> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + use tiberius::IntoRow; + + let table = format!("##{}", random_table().await); + let column_defs = &[$($sql_type,)+]; + + conn.execute( + &format!( + "CREATE TABLE {} (id INT IDENTITY PRIMARY KEY, {})", + table, + column_defs.join(", "), + ), + &[], + ) + .await?; + + let mut count = 0; + + $( + let mut req = conn.bulk_insert_columns(&table, $cols).await?; + for i in $generator { + let row = i.into_row(); + req.send(row).await?; + } + + let res = req.finalize().await?; + count += res.total(); + )+ + assert_eq!($total_generated, count); + + Ok(()) + } + + #[test_on_runtimes] + async fn [< bulk_load_required_ $name >](mut conn: tiberius::Client) -> Result<()> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + use tiberius::IntoRow; + let table = format!("##{}", random_table().await); + let column_defs = &[$(format!("{} NOT NULL", $sql_type),)+]; + + conn.execute( + &format!( + "CREATE TABLE {} (id INT IDENTITY PRIMARY KEY, {})", + table, + column_defs.join(", "), + ), + &[], + ) + .await?; + + let mut count = 0; + + $( + let mut req = conn.bulk_insert_columns(&table, $cols).await?; + for i in $generator { + let row = i.into_row(); + req.send(row).await?; + } + + let res = req.finalize().await?; + count += res.total(); + )+ + assert_eq!($total_generated, count); + + Ok(()) + } + + } + }; +} + +test_bulk_columns!(ab_ba_default_columns( + 200, + "a INT", + "b FLOAT", + "c INT DEFAULT 0", + (&["a", "b"], vec![(1i32, 1f64); 100]), + (&["b", "a"], vec![(2f64, 2i32); 100]), +)); + +test_bulk_columns!(ab_ba_override_default_columns( + 200, + "a INT", + "b FLOAT", + "c INT DEFAULT 0", + (&["a", "b", "c"], vec![(1i32, 1f64, 10i32); 100]), + (&["b", "c", "a"], vec![(2f64, 20i32, 2i32); 100]), +)); diff --git a/tests/query.rs b/tests/query.rs index 0a7b120e4..0cea71ebe 100644 --- a/tests/query.rs +++ b/tests/query.rs @@ -400,6 +400,33 @@ where Ok(()) } +#[test_on_runtimes] +async fn read_and_write_to_keyword_columns(mut conn: tiberius::Client) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let table = format!("##{}", random_table().await); + + conn.simple_query(format!("CREATE TABLE {} ([End] INT)", table)) + .await?; + + let res = conn + .execute(format!("INSERT INTO {} ([End]) VALUES (5)", table), &[]) + .await?; + + assert_eq!(1, res.total()); + + let rows = conn + .query(format!("SELECT [End] FROM {}", table), &[]) + .await? + .into_first_result() + .await?; + + assert_eq!(Some(5), rows[0].get(0)); + + Ok(()) +} + #[test_on_runtimes] async fn execute_insert_update_delete(mut conn: tiberius::Client) -> Result<()> where