tiberius/tds/codec/token/
token_done.rs

1use crate::{tds::codec::Encode, Error, SqlReadBytes, TokenType};
2use asynchronous_codec::BytesMut;
3use bytes::BufMut;
4use enumflags2::{bitflags, BitFlags};
5use std::fmt;
6
7#[derive(Debug, Default)]
8pub struct TokenDone {
9    status: BitFlags<DoneStatus>,
10    cur_cmd: u16,
11    done_rows: u64,
12}
13
14#[bitflags]
15#[repr(u16)]
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum DoneStatus {
18    More = 1 << 0,
19    Error = 1 << 1,
20    Inexact = 1 << 2,
21    // reserved
22    Count = 1 << 4,
23    Attention = 1 << 5,
24    // reserved
25    RpcInBatch = 1 << 7,
26    SrvError = 1 << 8,
27}
28
29impl TokenDone {
30    pub(crate) async fn decode<R>(src: &mut R) -> crate::Result<Self>
31    where
32        R: SqlReadBytes + Unpin,
33    {
34        let status = BitFlags::from_bits(src.read_u16_le().await?)
35            .map_err(|_| Error::Protocol("done(variant): invalid status".into()))?;
36
37        let cur_cmd = src.read_u16_le().await?;
38        let done_row_count_bytes = src.context().version().done_row_count_bytes();
39
40        let done_rows = match done_row_count_bytes {
41            8 => src.read_u64_le().await?,
42            4 => src.read_u32_le().await? as u64,
43            _ => unreachable!(),
44        };
45
46        Ok(TokenDone {
47            status,
48            cur_cmd,
49            done_rows,
50        })
51    }
52
53    pub(crate) fn is_final(&self) -> bool {
54        self.status.is_empty()
55    }
56
57    pub(crate) fn rows(&self) -> u64 {
58        self.done_rows
59    }
60}
61
62impl Encode<BytesMut> for TokenDone {
63    fn encode(self, dst: &mut BytesMut) -> crate::Result<()> {
64        dst.put_u8(TokenType::Done as u8);
65        dst.put_u16_le(BitFlags::bits(self.status));
66
67        dst.put_u16_le(self.cur_cmd);
68        dst.put_u64_le(self.done_rows);
69
70        Ok(())
71    }
72}
73
74impl fmt::Display for TokenDone {
75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
76        if self.done_rows == 0 {
77            write!(f, "Done with status {:?}", self.status)
78        } else if self.done_rows == 1 {
79            write!(f, "Done with status {:?} (1 row left)", self.status)
80        } else {
81            write!(
82                f,
83                "Done with status {:?} ({} rows left)",
84                self.status, self.done_rows
85            )
86        }
87    }
88}