Add method to query column metadata

src/client.rs

@@ -15,14 +15,10 @@ pub use config::*;
 pub(crate) use connection::*;
 
 use crate::tds::stream::ReceivedToken;
-use crate::{
-    result::ExecuteResult,
-    tds::{
-        codec::{self, IteratorJoin},
-        stream::{QueryStream, TokenStream},
-    },
-    BulkLoadRequest, ColumnFlag, SqlReadBytes, ToSql,
-};
+use crate::{result::ExecuteResult, tds::{
+    codec::{self, IteratorJoin},
+    stream::{QueryStream, TokenStream},
+}, BulkLoadRequest, ColumnFlag, MetaDataColumn, SqlReadBytes, ToSql};
 use codec::{BatchRequest, ColumnData, PacketHeader, RpcParam, RpcProcId, TokenRpcRequest};
 use enumflags2::BitFlags;
 use futures_util::io::{AsyncRead, AsyncWrite};
@@ -251,10 +247,13 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
         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
     ///
     /// ```
@@ -298,13 +297,96 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
     /// ```
     pub async fn bulk_insert<'a>(
         &'a mut self,
-        table: &'a str,
+        table: &str,
     ) -> crate::Result<BulkLoadRequest<'a, S>> {
+        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<dyn std::error::Error>> {
+    /// # 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: &str,
+        columns: &[&str],
+    ) -> crate::Result<BulkLoadRequest<'a, S>> {
+        let columns: Vec<MetaDataColumn<'a>> = self.column_metadata(table, columns).await?
+            .into_iter()
+            .filter(|column| column.base.flags.contains(ColumnFlag::Updateable))
+            .collect();
+
         // Start the bulk request
         self.connection.flush_stream().await?;
 
+        let col_data = columns.iter().map(MetaDataColumn::to_string).join(", ");
+        let query = format!("INSERT BULK {} ({})", table, col_data);
+
+        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 ts = TokenStream::new(&mut self.connection);
+        ts.flush_done().await?;
+
+        BulkLoadRequest::new(&mut self.connection, columns)
+    }
+
+    /// Retrieve the column metadata for a table, including column names, types,
+    /// sizes, and flags (e.g. nullability).
+    pub async fn column_metadata<'a, 'b>(
+        &'a mut self,
+        table: &str,
+        columns: &[&str],
+    ) -> crate::Result<Vec<MetaDataColumn<'b>>> {
+        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());
 
@@ -321,30 +403,12 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
 
                 Ok(columns)
             })
-            .await?;
-
-        // now start bulk upload
-        let columns: Vec<_> = columns
+            .await?
             .ok_or_else(|| {
                 crate::Error::Protocol("expecting column metadata from query but not found".into())
-            })?
-            .into_iter()
-            .filter(|column| column.base.flags.contains(ColumnFlag::Updateable))
-            .collect();
-
-        self.connection.flush_stream().await?;
-        let col_data = columns.iter().map(|c| format!("{}", c)).join(", ");
-        let query = format!("INSERT BULK {} ({})", table, col_data);
-
-        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 ts = TokenStream::new(&mut self.connection);
-        ts.flush_done().await?;
-
-        BulkLoadRequest::new(&mut self.connection, columns)
+        Ok(columns)
     }
 
     /// Closes this database connection explicitly.
@@ -371,7 +435,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
         &'a mut self,
         proc_id: RpcProcId,
         mut rpc_params: Vec<RpcParam<'b>>,
-        params: impl Iterator<Item = ColumnData<'b>>,
+        params: impl Iterator<Item=ColumnData<'b>>,
     ) -> crate::Result<()>
     where
         'a: 'b,

src/lib.rs

@@ -277,8 +277,9 @@ pub use result::*;
 pub use row::{Column, ColumnType, Row};
 pub use sql_browser::SqlBrowser;
 pub use tds::{
-    codec::{BulkLoadRequest, ColumnData, ColumnFlag, IntoRow, TokenRow, TypeLength},
+    codec::{BulkLoadRequest, ColumnData, MetaDataColumn, BaseMetaDataColumn, TypeInfo, FixedLenType, VarLenContext, VarLenType, ColumnFlag, IntoRow, TokenRow, TypeLength},
     numeric,
+    collation::{Collation},
     stream::QueryStream,
     time, xml, EncryptionLevel,
 };

src/row.rs

@@ -317,6 +317,11 @@ impl Row {
         self.result_index
     }
 
+    /// The raw column data
+    pub fn data(&self) -> &TokenRow<'static> {
+        &self.data
+    }
+
     /// Returns the number of columns in the row.
     ///
     /// # Example

src/tds.rs

@@ -1,5 +1,5 @@
 pub mod codec;
-mod collation;
+pub(crate) mod collation;
 mod context;
 pub mod numeric;
 pub mod stream;

src/tds/codec/token/token_col_metadata.rs

@@ -25,9 +25,14 @@ 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, self.base.ty)?;
 
-        match &self.base.ty {
+        Ok(())
+    }
+}
+impl Display for TypeInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self {
             TypeInfo::FixedLen(fixed) => match fixed {
                 FixedLenType::Int1 => write!(f, "tinyint")?,
                 FixedLenType::Bit => write!(f, "bit")?,
@@ -265,9 +270,9 @@ pub enum ColumnFlag {
     /// If column is writeable.
     Updateable = 1 << 3,
     /// Column modification status unknown.
-    UpdateableUnknown = 1 << 4,
+    UpdateableUnknown = 1 << 2,
     /// Column is an identity.
-    Identity = 1 << 5,
+    Identity = 1 << 4,
     /// Coulumn is computed.
     Computed = 1 << 7,
     /// Column is a fixed-length common language runtime user-defined type (CLR

src/tds/codec/token/token_row.rs

@@ -8,7 +8,7 @@ use futures_util::io::AsyncReadExt;
 pub use into_row::IntoRow;
 
 /// A row of data.
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, PartialEq)]
 pub struct TokenRow<'a> {
     data: Vec<ColumnData<'a>>,
 }

src/tds/numeric.rs

@@ -263,6 +263,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")]

tests/bulk.rs

@@ -218,3 +218,101 @@ test_bulk_type!(datetime2_7(
     100,
     vec![DateTime::from_timestamp(1658524194, 123456789); 100].into_iter()
 ));
+
+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 >]<S>(mut conn: tiberius::Client<S>) -> 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 >]<S>(mut conn: tiberius::Client<S>) -> 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]),
+));