mod bytes_mut_with_data_columns;
mod into_row;
use crate::tds::codec::encode::Encode;
use crate::{tds::codec::ColumnData, BytesMutWithTypeInfo, SqlReadBytes, TokenType};
use bytes::BufMut;
pub(crate) use bytes_mut_with_data_columns::BytesMutWithDataColumns;
use futures::io::AsyncReadExt;
pub use into_row::IntoRow;
#[derive(Debug, Default, Clone)]
pub struct TokenRow<'a> {
data: Vec<ColumnData<'a>>,
}
impl<'a> IntoIterator for TokenRow<'a> {
type Item = ColumnData<'a>;
type IntoIter = std::vec::IntoIter<Self::Item>;
fn into_iter(self) -> Self::IntoIter {
self.data.into_iter()
}
}
impl<'a> Encode<BytesMutWithDataColumns<'a>> for TokenRow<'a> {
fn encode(self, dst: &mut BytesMutWithDataColumns<'a>) -> crate::Result<()> {
dst.put_u8(TokenType::Row as u8);
if self.data.len() != dst.data_columns().len() {
return Err(crate::Error::BulkInput(
format!(
"Expecting {} columns but {} were given",
dst.data_columns().len(),
self.data.len()
)
.into(),
));
}
for (value, column) in self.data.into_iter().zip(dst.data_columns()) {
let mut dst_ti = BytesMutWithTypeInfo::new(dst).with_type_info(&column.base.ty);
value.encode(&mut dst_ti)?
}
Ok(())
}
}
impl<'a> TokenRow<'a> {
pub const fn new() -> Self {
Self { data: Vec::new() }
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
data: Vec::with_capacity(capacity),
}
}
pub fn clear(&mut self) {
self.data.clear();
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn get(&self, index: usize) -> Option<&ColumnData<'a>> {
self.data.get(index)
}
pub fn push(&mut self, value: ColumnData<'a>) {
self.data.push(value);
}
}
impl TokenRow<'static> {
pub(crate) async fn decode<R>(src: &mut R) -> crate::Result<Self>
where
R: SqlReadBytes + Unpin,
{
let col_meta = src.context().last_meta().unwrap();
let mut row = Self {
data: Vec::with_capacity(col_meta.columns.len()),
};
for column in col_meta.columns.iter() {
let data = ColumnData::decode(src, &column.base.ty).await?;
row.data.push(data);
}
Ok(row)
}
pub(crate) async fn decode_nbc<R>(src: &mut R) -> crate::Result<Self>
where
R: SqlReadBytes + Unpin,
{
let col_meta = src.context().last_meta().unwrap();
let row_bitmap = RowBitmap::decode(src, col_meta.columns.len()).await?;
let mut row = Self {
data: Vec::with_capacity(col_meta.columns.len()),
};
for (i, column) in col_meta.columns.iter().enumerate() {
let data = if row_bitmap.is_null(i) {
column.base.null_value()
} else {
ColumnData::decode(src, &column.base.ty).await?
};
row.data.push(data);
}
Ok(row)
}
}
struct RowBitmap {
data: Vec<u8>,
}
impl RowBitmap {
#[inline]
fn is_null(&self, i: usize) -> bool {
let index = i / 8;
let bit = i % 8;
self.data[index] & (1 << bit) > 0
}
async fn decode<R>(src: &mut R, columns: usize) -> crate::Result<Self>
where
R: SqlReadBytes + Unpin,
{
let size = (columns + 8 - 1) / 8;
let mut data = vec![0; size];
src.read_exact(&mut data[0..size]).await?;
Ok(Self { data })
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{BaseMetaDataColumn, ColumnFlag, FixedLenType, MetaDataColumn, TypeInfo};
use bytes::BytesMut;
#[tokio::test]
async fn wrong_number_of_columns_will_fail() {
let row = (true, 5).into_row();
let columns = vec![MetaDataColumn {
base: BaseMetaDataColumn {
flags: ColumnFlag::Nullable.into(),
ty: TypeInfo::FixedLen(FixedLenType::Bit),
},
col_name: Default::default(),
}];
let mut buf = BytesMut::new();
let mut buf_with_columns = BytesMutWithDataColumns::new(&mut buf, &columns);
row.encode(&mut buf_with_columns)
.expect_err("wrong number of columns");
}
}