1use std::sync::LazyLock;
16
17use anyhow::{Context, bail, ensure};
18use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
19use risingwave_common::catalog::Schema;
20use risingwave_common::log::LogSuppressor;
21use risingwave_common::row::OwnedRow;
22use risingwave_common::types::{
23 DataType, Date, Datum, Decimal, ScalarImpl, Time, Timestamp, Timestamptz,
24};
25use risingwave_common::util::iter_util::ZipEqFast;
26use rust_decimal::Decimal as RustDecimal;
27use thiserror_ext::AsReport;
28use tiberius::xml::XmlData;
29use tiberius::{Column, ColumnType, Row};
30use uuid::Uuid;
31
32use crate::parser::utils::log_error;
33
34static LOG_SUPPRESSOR: LazyLock<LogSuppressor> = LazyLock::new(LogSuppressor::default);
35
36pub struct SqlServerRowDecoder {
43 schema: Schema,
44 money_columns: Vec<bool>,
45}
46
47impl SqlServerRowDecoder {
48 pub fn new(columns: &[Column], schema: &Schema) -> anyhow::Result<Self> {
49 ensure!(
50 columns.len() == schema.len(),
51 "SQL Server returned {} columns, but the planned schema has {} columns",
52 columns.len(),
53 schema.len()
54 );
55
56 let money_columns = columns
57 .iter()
58 .zip_eq_fast(&schema.fields)
59 .map(|(column, field)| {
60 let column_type = column.column_type();
61 match (column_type, &field.data_type) {
62 (ColumnType::Money | ColumnType::Money4, DataType::Decimal) => Ok(true),
63 (ColumnType::Money | ColumnType::Money4, data_type) => bail!(
64 "SQL Server MONEY column `{}` has incompatible RisingWave type {}",
65 field.name,
66 data_type
67 ),
68 (ColumnType::Decimaln | ColumnType::Numericn, DataType::Decimal) => Ok(false),
69 (_, DataType::Decimal) => bail!(
70 "SQL Server column `{}` has TDS type {:?}, which is incompatible with the planned RisingWave Decimal type",
71 field.name,
72 column_type
73 ),
74 _ => Ok(false),
75 }
76 })
77 .collect::<anyhow::Result<Vec<_>>>()?;
78
79 Ok(Self {
80 schema: schema.clone(),
81 money_columns,
82 })
83 }
84
85 fn validate_row(&self, row: &Row) -> anyhow::Result<()> {
86 ensure!(
87 row.result_index() == 0,
88 "mssql_query returned more than one result set"
89 );
90 ensure!(
91 row.columns().len() == self.schema.len(),
92 "SQL Server row has {} columns, but the planned schema has {} columns",
93 row.columns().len(),
94 self.schema.len()
95 );
96 Ok(())
97 }
98
99 pub fn decode(&self, row: &mut Row) -> anyhow::Result<OwnedRow> {
100 self.validate_row(row)?;
101
102 let mut datums = Vec::with_capacity(self.schema.len());
103 for (i, rw_field) in self.schema.fields.iter().enumerate() {
104 let name = rw_field.name.as_str();
105 let datum = match sql_server_cell_to_rw_datum(
106 row,
107 i,
108 name,
109 &rw_field.data_type,
110 self.money_columns[i],
111 ) {
112 Ok(datum) => datum,
113 Err(err) => {
114 log_error!(name, err, "parse column failed");
115 None
116 }
117 };
118 datums.push(datum);
119 }
120 Ok(OwnedRow::new(datums))
121 }
122
123 pub fn decode_with_strict_pk(
126 &self,
127 row: &mut Row,
128 pk_indices: &[usize],
129 ) -> anyhow::Result<OwnedRow> {
130 self.validate_row(row)?;
131
132 super::decode_row_with_strict_pk(
133 "SQL Server",
134 &self.schema,
135 pk_indices,
136 |index, field| {
137 sql_server_cell_to_rw_datum(
138 row,
139 index,
140 &field.name,
141 &field.data_type,
142 self.money_columns[index],
143 )
144 },
145 |name, err| log_error!(name, err, "parse column failed"),
146 )
147 }
148}
149
150fn sql_server_cell_to_rw_datum(
151 row: &Row,
152 index: usize,
153 name: &str,
154 data_type: &DataType,
155 is_money: bool,
156) -> anyhow::Result<Datum> {
157 if is_money {
158 return row
159 .try_get::<i64, usize>(index)
160 .with_context(|| format!("failed to decode SQL Server money column `{name}`"))?
161 .map(|value| try_convert_money_i64_to_type(value, data_type))
162 .transpose();
163 }
164
165 Ok(row
166 .try_get::<ScalarImplTiberiusWrapper, usize>(index)
167 .with_context(|| format!("failed to decode SQL Server snapshot column `{name}`"))?
168 .map(|datum| datum.0)
169 .map(|scalar| coerce_scalar_to_target_type(scalar, data_type)))
170}
171
172fn coerce_scalar_to_target_type(scalar: ScalarImpl, target_type: &DataType) -> ScalarImpl {
173 match (scalar, target_type) {
174 (ScalarImpl::Int16(v), DataType::Int32) => ScalarImpl::Int32(v as i32),
177 (ScalarImpl::Int16(v), DataType::Int64) => ScalarImpl::Int64(v as i64),
178 (ScalarImpl::Int32(v), DataType::Int64) => ScalarImpl::Int64(v as i64),
179 (ScalarImpl::Float32(v), DataType::Float64) => ScalarImpl::Float64((v.0 as f64).into()),
181 (scalar, _) => scalar,
182 }
183}
184
185fn try_convert_money_i64_to_type(value: i64, data_type: &DataType) -> anyhow::Result<ScalarImpl> {
186 match data_type {
187 DataType::Decimal => {
188 Ok(ScalarImpl::Decimal(Decimal::from_i128_with_scale(
190 value as i128,
191 4,
192 )))
193 }
194 _ => bail!("conversion of SQL Server money to {data_type} is not supported"),
195 }
196}
197
198#[cfg(test)]
199mod tests {
200 use risingwave_common::catalog::Field;
201 use risingwave_common::types::F32;
202
203 use super::*;
204
205 #[test]
206 fn test_row_decoder_uses_tds_money_metadata() {
207 let columns = vec![
208 Column::new("money".into(), ColumnType::Money),
209 Column::new("decimal".into(), ColumnType::Decimaln),
210 ];
211 let schema = Schema::new(vec![
212 Field::with_name(DataType::Decimal, "money"),
213 Field::with_name(DataType::Decimal, "decimal"),
214 ]);
215
216 let decoder = SqlServerRowDecoder::new(&columns, &schema).unwrap();
217 assert_eq!(decoder.money_columns, vec![true, false]);
218
219 assert!(SqlServerRowDecoder::new(&columns[..1], &schema).is_err());
220
221 let intn_column = [Column::new("not_money".into(), ColumnType::Intn)];
222 let decimal_schema = Schema::new(vec![Field::with_name(DataType::Decimal, "not_money")]);
223 assert!(SqlServerRowDecoder::new(&intn_column, &decimal_schema).is_err());
224 }
225
226 #[test]
227 fn test_integer_upcast_coercion() {
228 let v = coerce_scalar_to_target_type(ScalarImpl::Int32(7), &DataType::Int64);
229 assert_eq!(v, ScalarImpl::Int64(7));
230
231 let v = coerce_scalar_to_target_type(ScalarImpl::Int16(7), &DataType::Int32);
232 assert_eq!(v, ScalarImpl::Int32(7));
233
234 let v = coerce_scalar_to_target_type(ScalarImpl::Int16(7), &DataType::Int64);
235 assert_eq!(v, ScalarImpl::Int64(7));
236 }
237
238 #[test]
239 fn test_float_upcast_coercion() {
240 let v =
241 coerce_scalar_to_target_type(ScalarImpl::Float32(F32::from(1.25)), &DataType::Float64);
242 assert_eq!(v, ScalarImpl::Float64(1.25.into()));
243 }
244
245 #[test]
246 fn test_non_upcast_keeps_original() {
247 let v = coerce_scalar_to_target_type(ScalarImpl::Int32(7), &DataType::Int32);
248 assert_eq!(v, ScalarImpl::Int32(7));
249 }
250
251 #[test]
255 fn test_money_i64_to_decimal() {
256 fn assert_money(raw: i64, expected: Decimal) {
257 let ScalarImpl::Decimal(actual) =
258 try_convert_money_i64_to_type(raw, &DataType::Decimal).unwrap()
259 else {
260 unreachable!()
261 };
262 assert_eq!(actual, expected);
263 assert_eq!(actual.scale(), Some(4));
264 }
265
266 assert_money(10000, Decimal::from(1));
267
268 assert_money(12300, Decimal::from(123) / Decimal::from(100));
270
271 assert_money(12345600, Decimal::from(123456) / Decimal::from(100));
273
274 assert_money(0, Decimal::from(0));
276 }
277
278 #[test]
282 fn test_money_i64_to_unsupported_type_errors() {
283 for unsupported in [
284 DataType::Int32,
285 DataType::Int64,
286 DataType::Float64,
287 DataType::Varchar,
288 ] {
289 assert!(
290 try_convert_money_i64_to_type(12345600, &unsupported).is_err(),
291 "expected error for money -> {unsupported:?}"
292 );
293 }
294 }
295}
296macro_rules! impl_tiberius_wrapper {
297 ($wrapper_name:ident, $variant_name:ident) => {
298 pub struct $wrapper_name($variant_name);
299
300 impl From<$variant_name> for $wrapper_name {
301 fn from(value: $variant_name) -> Self {
302 Self(value)
303 }
304 }
305 };
306}
307
308impl_tiberius_wrapper!(ScalarImplTiberiusWrapper, ScalarImpl);
309impl_tiberius_wrapper!(TimeTiberiusWrapper, Time);
310impl_tiberius_wrapper!(DateTiberiusWrapper, Date);
311impl_tiberius_wrapper!(TimestampTiberiusWrapper, Timestamp);
312impl_tiberius_wrapper!(TimestamptzTiberiusWrapper, Timestamptz);
313impl_tiberius_wrapper!(DecimalTiberiusWrapper, Decimal);
314
315macro_rules! impl_chrono_tiberius_wrapper {
316 ($wrapper_name:ident, $variant_name:ident, $chrono:ty) => {
317 impl<'a> tiberius::IntoSql<'a> for $wrapper_name {
318 fn into_sql(self) -> tiberius::ColumnData<'a> {
319 self.0.0.into_sql()
320 }
321 }
322
323 impl<'a> tiberius::FromSql<'a> for $wrapper_name {
324 fn from_sql(
325 value: &'a tiberius::ColumnData<'static>,
326 ) -> tiberius::Result<Option<Self>> {
327 let instant = <$chrono>::from_sql(value)?;
328 let time = instant.map($variant_name::from).map($wrapper_name::from);
329 tiberius::Result::Ok(time)
330 }
331 }
332 };
333}
334
335impl_chrono_tiberius_wrapper!(TimeTiberiusWrapper, Time, NaiveTime);
336impl_chrono_tiberius_wrapper!(DateTiberiusWrapper, Date, NaiveDate);
337impl_chrono_tiberius_wrapper!(TimestampTiberiusWrapper, Timestamp, NaiveDateTime);
338
339impl<'a> tiberius::IntoSql<'a> for DecimalTiberiusWrapper {
340 fn into_sql(self) -> tiberius::ColumnData<'a> {
341 match self.0 {
342 Decimal::Normalized(d) => d.into_sql(),
343 Decimal::NaN => tiberius::ColumnData::Numeric(None),
344 Decimal::PositiveInf => tiberius::ColumnData::Numeric(None),
345 Decimal::NegativeInf => tiberius::ColumnData::Numeric(None),
346 }
347 }
348}
349
350impl<'a> tiberius::FromSql<'a> for DecimalTiberiusWrapper {
351 fn from_sql(value: &'a tiberius::ColumnData<'static>) -> tiberius::Result<Option<Self>> {
353 tiberius::Result::Ok(
354 RustDecimal::from_sql(value)?
355 .map(Decimal::Normalized)
356 .map(DecimalTiberiusWrapper::from),
357 )
358 }
359}
360
361impl<'a> tiberius::IntoSql<'a> for TimestamptzTiberiusWrapper {
362 fn into_sql(self) -> tiberius::ColumnData<'a> {
363 self.0.to_datetime_utc().into_sql()
364 }
365}
366
367impl<'a> tiberius::FromSql<'a> for TimestamptzTiberiusWrapper {
368 fn from_sql(value: &'a tiberius::ColumnData<'static>) -> tiberius::Result<Option<Self>> {
369 let instant = time::OffsetDateTime::from_sql(value)?;
370 instant
371 .map(|instant| {
372 let timestamptz = instant
373 .unix_timestamp_nanos()
374 .checked_div(1000)
375 .and_then(|micros| i64::try_from(micros).ok())
376 .and_then(Timestamptz::from_micros)
377 .ok_or_else(|| {
378 tiberius::error::Error::Conversion(
379 "datetimeoffset is out of range for RisingWave timestamptz".into(),
380 )
381 })?;
382 Ok(TimestamptzTiberiusWrapper::from(timestamptz))
383 })
384 .transpose()
385 }
386}
387
388impl<'a> tiberius::FromSql<'a> for ScalarImplTiberiusWrapper {
422 fn from_sql(value: &'a tiberius::ColumnData<'static>) -> tiberius::Result<Option<Self>> {
423 Ok(match &value {
424 tiberius::ColumnData::U8(_) => u8::from_sql(value)?
425 .map(|v| ScalarImplTiberiusWrapper::from(ScalarImpl::from(v as i16))),
426 tiberius::ColumnData::I16(_) => i16::from_sql(value)?
427 .map(ScalarImpl::from)
428 .map(ScalarImplTiberiusWrapper::from),
429 tiberius::ColumnData::I32(_) => i32::from_sql(value)?
430 .map(ScalarImpl::from)
431 .map(ScalarImplTiberiusWrapper::from),
432 tiberius::ColumnData::I64(_) => i64::from_sql(value)?
433 .map(ScalarImpl::from)
434 .map(ScalarImplTiberiusWrapper::from),
435 tiberius::ColumnData::F32(_) => f32::from_sql(value)?
436 .map(ScalarImpl::from)
437 .map(ScalarImplTiberiusWrapper::from),
438 tiberius::ColumnData::F64(_) => f64::from_sql(value)?
439 .map(ScalarImpl::from)
440 .map(ScalarImplTiberiusWrapper::from),
441 tiberius::ColumnData::Bit(_) => bool::from_sql(value)?
442 .map(ScalarImpl::from)
443 .map(ScalarImplTiberiusWrapper::from),
444 tiberius::ColumnData::String(_) => <&str>::from_sql(value)?
445 .map(ScalarImpl::from)
446 .map(ScalarImplTiberiusWrapper::from),
447 tiberius::ColumnData::Numeric(_) => DecimalTiberiusWrapper::from_sql(value)?
448 .map(|w| ScalarImpl::from(w.0))
449 .map(ScalarImplTiberiusWrapper::from),
450 tiberius::ColumnData::DateTime(_)
451 | tiberius::ColumnData::DateTime2(_)
452 | tiberius::ColumnData::SmallDateTime(_) => TimestampTiberiusWrapper::from_sql(value)?
453 .map(|w| ScalarImpl::from(w.0))
454 .map(ScalarImplTiberiusWrapper::from),
455 tiberius::ColumnData::Time(_) => TimeTiberiusWrapper::from_sql(value)?
456 .map(|w| ScalarImpl::from(w.0))
457 .map(ScalarImplTiberiusWrapper::from),
458 tiberius::ColumnData::Date(_) => DateTiberiusWrapper::from_sql(value)?
459 .map(|w| ScalarImpl::from(w.0))
460 .map(ScalarImplTiberiusWrapper::from),
461 tiberius::ColumnData::DateTimeOffset(_) => TimestamptzTiberiusWrapper::from_sql(value)?
462 .map(|w| ScalarImpl::from(w.0))
463 .map(ScalarImplTiberiusWrapper::from),
464 tiberius::ColumnData::Binary(_) => <&[u8]>::from_sql(value)?
465 .map(ScalarImpl::from)
466 .map(ScalarImplTiberiusWrapper::from),
467 tiberius::ColumnData::Guid(_) => <Uuid>::from_sql(value)?
468 .map(|uuid| uuid.to_string().to_uppercase())
469 .map(ScalarImpl::from)
470 .map(ScalarImplTiberiusWrapper::from),
471 tiberius::ColumnData::Xml(_) => <&XmlData>::from_sql(value)?
472 .map(|xml| xml.clone().into_string())
473 .map(ScalarImpl::from)
474 .map(ScalarImplTiberiusWrapper::from),
475 })
476 }
477}
478
479impl<'a> tiberius::IntoSql<'a> for ScalarImplTiberiusWrapper {
527 fn into_sql(self) -> tiberius::ColumnData<'a> {
528 match self.0 {
529 ScalarImpl::Int16(v) => v.into_sql(),
530 ScalarImpl::Int32(v) => v.into_sql(),
531 ScalarImpl::Int64(v) => v.into_sql(),
532 ScalarImpl::Float32(v) => v.0.into_sql(),
533 ScalarImpl::Float64(v) => v.0.into_sql(),
534 ScalarImpl::Bool(v) => v.into_sql(),
535 ScalarImpl::Decimal(v) => DecimalTiberiusWrapper::from(v).into_sql(),
536 ScalarImpl::Date(v) => DateTiberiusWrapper::from(v).into_sql(),
537 ScalarImpl::Timestamp(v) => TimestampTiberiusWrapper::from(v).into_sql(),
538 ScalarImpl::Timestamptz(v) => TimestamptzTiberiusWrapper::from(v).into_sql(),
539 ScalarImpl::Time(v) => TimeTiberiusWrapper::from(v).into_sql(),
540 ScalarImpl::Bytea(v) => {
541 let value: Vec<u8> = (*v).to_vec();
542 value.into_sql()
543 }
544 ScalarImpl::Utf8(v) => String::from(v).into_sql(),
545 value => {
546 unimplemented!("the sql server decoding for {:?} is unsupported", value);
548 }
549 }
550 }
551}