1mod arrow_convert;
19mod proxy;
20
21pub use proxy::MssqlProxyKind;
22
23use std::sync::Arc;
24
25use arrow::datatypes::SchemaRef;
26use tiberius::{AuthMethod, Client, Config, EncryptionLevel};
27use tokio::net::TcpStream;
28use tokio::runtime::Runtime;
29use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
30
31use proxy::{detect_mssql_proxy_kind, warn_proxy_kind};
32
33use crate::config::TlsConfig;
34use crate::error::Result;
35use crate::source::batch_controller::{
36 AdaptiveBatchController, DEFAULT_BATCH_TARGET_MB, PROBE_BATCH_SIZE,
37};
38use crate::source::query::build_export_query;
39use crate::source::{BatchSink, ExportRequest, Source, TableIntrospection};
40use crate::types::{ColumnOverrides, TypeMapping};
41
42type MssqlClient = Client<Compat<TcpStream>>;
43
44pub struct MssqlSource {
50 rt: Runtime,
51 client: MssqlClient,
52 proxy_kind: MssqlProxyKind,
54 lock_timeout_applied: bool,
57}
58
59impl Drop for MssqlSource {
60 fn drop(&mut self) {
74 if !self.lock_timeout_applied {
75 return;
76 }
77 let Self { rt, client, .. } = self;
78 let _ = rt.block_on(async {
79 tokio::time::timeout(
80 std::time::Duration::from_secs(2),
81 client.execute("SET LOCK_TIMEOUT -1", &[]),
82 )
83 .await
84 });
85 }
86}
87
88struct MssqlUrl {
90 host: String,
91 port: u16,
92 user: String,
93 password: String,
94 database: String,
95}
96
97fn parse_mssql_url(url: &str) -> Result<MssqlUrl> {
98 let rest = url
99 .strip_prefix("sqlserver://")
100 .or_else(|| url.strip_prefix("mssql://"))
101 .ok_or_else(|| anyhow::anyhow!("mssql url must start with sqlserver:// — got {url}"))?;
102 let (userinfo, hostpart) = rest
105 .rsplit_once('@')
106 .ok_or_else(|| anyhow::anyhow!("mssql url missing user@host: {url}"))?;
107 let (user, password) = match userinfo.split_once(':') {
108 Some((u, p)) => (u.to_string(), p.to_string()),
109 None => (userinfo.to_string(), String::new()),
110 };
111 let (hostport, database) = hostpart
112 .split_once('/')
113 .map(|(h, d)| (h, d.to_string()))
114 .unwrap_or((hostpart, String::new()));
115 let (host, port) = match hostport.rsplit_once(':') {
116 Some((h, p)) => (
117 h.to_string(),
118 p.parse::<u16>()
119 .map_err(|_| anyhow::anyhow!("mssql url port not a number: {p}"))?,
120 ),
121 None => (hostport.to_string(), 1433),
122 };
123 if database.is_empty() {
124 anyhow::bail!("mssql url must include a database: sqlserver://user:pass@host:port/<db>");
125 }
126 Ok(MssqlUrl {
127 host,
128 port,
129 user,
130 password,
131 database,
132 })
133}
134
135impl MssqlSource {
136 pub fn connect_with_tls(url: &str, tls: Option<&TlsConfig>) -> Result<Self> {
140 let parts = parse_mssql_url(url)?;
141 let mut config = Config::new();
142 config.host(&parts.host);
143 config.port(parts.port);
144 config.database(&parts.database);
145 config.authentication(AuthMethod::sql_server(&parts.user, &parts.password));
146
147 config.encryption(EncryptionLevel::Required);
152 match tls {
153 Some(cfg) if cfg.accept_invalid_certs => config.trust_cert(),
154 Some(cfg) => {
155 if let Some(ca) = cfg.ca_file.as_deref() {
156 config.trust_cert_ca(ca);
157 }
158 }
159 None => config.trust_cert(),
160 }
161
162 let rt = tokio::runtime::Builder::new_current_thread()
163 .enable_all()
164 .build()
165 .map_err(|e| anyhow::anyhow!("mssql: tokio runtime build failed: {e}"))?;
166
167 let client = rt.block_on(async {
168 let tcp = TcpStream::connect(config.get_addr())
169 .await
170 .map_err(|e| anyhow::anyhow!("mssql: TCP connect failed: {e}"))?;
171 tcp.set_nodelay(true).ok();
172 Client::connect(config, tcp.compat_write())
173 .await
174 .map_err(|e| anyhow::anyhow!("mssql: login failed: {e}"))
175 })?;
176
177 let mut src = Self {
178 rt,
179 client,
180 proxy_kind: MssqlProxyKind::Direct,
181 lock_timeout_applied: false,
182 };
183 src.query_scalar("SELECT 1")?;
186 let kind = detect_mssql_proxy_kind(&src.rt, &mut src.client);
191 warn_proxy_kind(kind);
192 src.proxy_kind = kind;
193 Ok(src)
194 }
195
196 #[allow(dead_code)]
200 pub fn proxy_kind(&self) -> MssqlProxyKind {
201 self.proxy_kind
202 }
203}
204
205impl Source for MssqlSource {
206 fn export(&mut self, request: &ExportRequest<'_>, sink: &mut dyn BatchSink) -> Result<()> {
207 let built = build_export_query(request, crate::config::SourceType::Mssql);
210 let sql = built.sql.clone();
211 let overrides = request.column_overrides.clone();
212 let mut ctl =
221 AdaptiveBatchController::new(request.tuning, request.tuning.batch_size.max(1));
222 let mut cap_applied = false;
223 let lock_timeout_ms = request.tuning.lock_timeout_s.saturating_mul(1000);
233 let stmt_timeout = (request.tuning.statement_timeout_s > 0)
234 .then(|| std::time::Duration::from_secs(request.tuning.statement_timeout_s));
235
236 if lock_timeout_ms > 0 {
239 self.lock_timeout_applied = true;
240 }
241
242 let Self { rt, client, .. } = self;
243 rt.block_on(async {
244 use futures_util::stream::TryStreamExt;
245 use tiberius::QueryItem;
246
247 if lock_timeout_ms > 0 {
248 client
249 .execute(format!("SET LOCK_TIMEOUT {lock_timeout_ms}"), &[])
250 .await
251 .map_err(|e| anyhow::anyhow!("mssql: SET LOCK_TIMEOUT failed: {e}"))?;
252 }
253
254 let started = std::time::Instant::now();
255 let mut stream = client
256 .query(sql.as_str(), &[])
257 .await
258 .map_err(|e| anyhow::anyhow!("mssql: query failed: {e}"))?;
259
260 let mut columns: Vec<tiberius::Column> = Vec::new();
261 let mut buf: Vec<tiberius::Row> = Vec::with_capacity(ctl.target());
262 let mut schema: Option<SchemaRef> = None;
263
264 while let Some(item) = stream
265 .try_next()
266 .await
267 .map_err(|e| anyhow::anyhow!("mssql: streaming rows failed: {e}"))?
268 {
269 if let Some(budget) = stmt_timeout
270 && started.elapsed() > budget
271 {
272 anyhow::bail!(
273 "mssql: statement timeout after {}s (tuning.statement_timeout_s)",
274 budget.as_secs()
275 );
276 }
277 match item {
278 QueryItem::Metadata(meta) if columns.is_empty() => {
281 columns = meta.columns().to_vec();
282 }
283 QueryItem::Metadata(_) => {}
284 QueryItem::Row(row) => {
285 buf.push(row);
286 if buf.len() >= ctl.target() {
287 let arrow_bytes =
288 emit_mssql_batch(&columns, &overrides, &mut schema, &buf, sink)?;
289 let n = buf.len();
290 buf.clear();
291 if !cap_applied && n > 0 {
296 let arrow_per_row = (arrow_bytes / n).max(64);
297 let target_mb = request
298 .tuning
299 .batch_size_memory_mb
300 .unwrap_or(DEFAULT_BATCH_TARGET_MB);
301 let safe = ((target_mb * 1024 * 1024) / arrow_per_row)
302 .max(PROBE_BATCH_SIZE);
303 if let Some(new) = ctl.apply_memory_cap(safe) {
304 log::info!(
305 "MSSQL batch cap: arrow≈{} B/row, target={} MB → batch_size → {}",
306 arrow_per_row,
307 target_mb,
308 new
309 );
310 buf.reserve(new.saturating_sub(buf.capacity()));
311 }
312 cap_applied = true;
313 }
314 ctl.after_batch(|| None);
316 ctl.throttle();
317 }
318 }
319 }
320 }
321 if !buf.is_empty() || schema.is_none() {
327 emit_mssql_batch(&columns, &overrides, &mut schema, &buf, sink)?;
328 }
329 Ok::<_, anyhow::Error>(())
330 })?;
331 Ok(())
332 }
333
334 fn query_scalar(&mut self, sql: &str) -> Result<Option<String>> {
335 let Self { rt, client, .. } = self;
336 rt.block_on(async {
337 let row = client
338 .query(sql, &[])
339 .await
340 .map_err(|e| anyhow::anyhow!("mssql: scalar query failed: {e}"))?
341 .into_row()
342 .await
343 .map_err(|e| anyhow::anyhow!("mssql: reading scalar row failed: {e}"))?;
344 Ok(row.and_then(|r| scalar_to_string(&r)))
345 })
346 }
347
348 fn type_mappings(
349 &mut self,
350 query: &str,
351 column_overrides: &ColumnOverrides,
352 ) -> Result<Vec<TypeMapping>> {
353 let wrapped = format!("SELECT * FROM ({query}) AS _rivet_q WHERE 1 = 0");
355 let overrides = column_overrides.clone();
356 let Self { rt, client, .. } = self;
357 rt.block_on(async {
358 let mut stream = client
359 .query(wrapped.as_str(), &[])
360 .await
361 .map_err(|e| anyhow::anyhow!("mssql: type-probe query failed: {e}"))?;
362 let columns = stream
363 .columns()
364 .await
365 .map_err(|e| anyhow::anyhow!("mssql: type-probe metadata failed: {e}"))?
366 .map(<[_]>::to_vec)
367 .unwrap_or_default();
368 let _ = stream.into_first_result().await;
370 Ok(arrow_convert::mssql_type_mappings(&columns, &overrides))
371 })
372 }
373
374 fn sample_pressure(&mut self) -> Option<u64> {
375 let Self { rt, client, .. } = self;
376 let sql = "SELECT SUM(cntr_value) FROM sys.dm_os_performance_counters \
386 WHERE counter_name IN ('Workfiles Created/sec', 'Worktables Created/sec')";
387 rt.block_on(async {
388 let row = client.query(sql, &[]).await.ok()?.into_row().await.ok()??;
389 row.get::<i64, _>(0).map(|v| v.max(0) as u64)
390 })
391 }
392}
393
394fn emit_mssql_batch(
404 columns: &[tiberius::Column],
405 overrides: &ColumnOverrides,
406 schema: &mut Option<SchemaRef>,
407 rows: &[tiberius::Row],
408 sink: &mut dyn BatchSink,
409) -> Result<usize> {
410 let schema_ref = match schema {
411 Some(s) => s.clone(),
412 None => {
413 let (built, _decoders) =
414 arrow_convert::mssql_columns_to_schema(columns, overrides, rows)?;
415 let s: SchemaRef = Arc::new(built);
416 sink.on_schema(s.clone())?;
417 *schema = Some(s.clone());
418 s
419 }
420 };
421 if !rows.is_empty() {
422 let batch = arrow_convert::mssql_rows_to_record_batch(&schema_ref, rows)?;
423 let bytes = crate::tuning::SourceTuning::batch_memory_bytes(&batch);
424 sink.on_batch(&batch)?;
425 return Ok(bytes);
426 }
427 Ok(0)
428}
429
430fn scalar_to_string(row: &tiberius::Row) -> Option<String> {
434 use tiberius::ColumnData;
435 let cell = row.cells().next().map(|(_, d)| d)?;
436 match cell {
437 ColumnData::U8(v) => v.map(|x| x.to_string()),
438 ColumnData::I16(v) => v.map(|x| x.to_string()),
439 ColumnData::I32(v) => v.map(|x| x.to_string()),
440 ColumnData::I64(v) => v.map(|x| x.to_string()),
441 ColumnData::F32(v) => v.map(|x| x.to_string()),
442 ColumnData::F64(v) => v.map(|x| x.to_string()),
443 ColumnData::Bit(v) => v.map(|x| x.to_string()),
444 ColumnData::String(v) => v.as_ref().map(|s| s.to_string()),
445 ColumnData::Numeric(v) => v.map(|n| {
446 let raw = n.value();
448 let scale = n.scale() as usize;
449 if scale == 0 {
450 raw.to_string()
451 } else {
452 let neg = raw < 0;
453 let digits = raw.unsigned_abs().to_string();
454 let digits = format!("{digits:0>width$}", width = scale + 1);
455 let (int, frac) = digits.split_at(digits.len() - scale);
456 format!("{}{int}.{frac}", if neg { "-" } else { "" })
457 }
458 }),
459 ColumnData::Guid(v) => v.map(|g| g.to_string()),
460 other => Some(format!("{other:?}")),
461 }
462}
463
464pub(crate) fn introspect_mssql_table_for_chunking(
467 url: &str,
468 tls: Option<&TlsConfig>,
469 qualified_table: &str,
470) -> Result<TableIntrospection> {
471 let (schema, table) = match qualified_table.split_once('.') {
472 Some((s, t)) => (s.to_string(), t.to_string()),
473 None => ("dbo".to_string(), qualified_table.to_string()),
474 };
475 let mut src = MssqlSource::connect_with_tls(url, tls)?;
476
477 let count_sql = format!(
480 "SELECT SUM(p.row_count) FROM sys.dm_db_partition_stats p \
481 JOIN sys.objects o ON o.object_id = p.object_id \
482 JOIN sys.schemas s ON s.schema_id = o.schema_id \
483 WHERE s.name = N'{}' AND o.name = N'{}' AND p.index_id IN (0,1)",
484 schema.replace('\'', "''"),
485 table.replace('\'', "''")
486 );
487 let row_estimate = src
488 .query_scalar(&count_sql)?
489 .and_then(|s| s.parse::<i64>().ok())
490 .unwrap_or(0);
491
492 let pk_sql = format!(
495 "SELECT TOP 1 c.name, t.name FROM sys.indexes i \
496 JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id \
497 JOIN sys.columns c ON c.object_id = ic.object_id AND c.column_id = ic.column_id \
498 JOIN sys.types t ON t.user_type_id = c.user_type_id \
499 JOIN sys.objects o ON o.object_id = i.object_id \
500 JOIN sys.schemas s ON s.schema_id = o.schema_id \
501 WHERE i.is_primary_key = 1 AND s.name = N'{}' AND o.name = N'{}' \
502 GROUP BY c.name, t.name HAVING COUNT(*) = 1",
503 schema.replace('\'', "''"),
504 table.replace('\'', "''")
505 );
506 let keyset_sql = format!(
516 "SELECT STRING_AGG(col, CHAR(31)) WITHIN GROUP (ORDER BY is_pk DESC, col) FROM ( \
517 SELECT col, MAX(is_pk) AS is_pk FROM ( \
518 SELECT MIN(c.name) AS col, MAX(CONVERT(int, i.is_primary_key)) AS is_pk \
519 FROM sys.indexes i \
520 JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id AND ic.key_ordinal > 0 \
521 JOIN sys.columns c ON c.object_id = ic.object_id AND c.column_id = ic.column_id \
522 JOIN sys.objects o ON o.object_id = i.object_id \
523 JOIN sys.schemas s ON s.schema_id = o.schema_id \
524 WHERE i.is_unique = 1 AND c.is_nullable = 0 AND s.name = N'{}' AND o.name = N'{}' \
525 GROUP BY i.object_id, i.index_id HAVING COUNT(*) = 1 \
526 ) per_index GROUP BY col \
527 ) deduped",
528 schema.replace('\'', "''"),
529 table.replace('\'', "''")
530 );
531 let keyset_keys: Vec<String> = src
532 .query_scalar(&keyset_sql)?
533 .map(|s| {
534 s.split('\u{1f}')
535 .filter(|c| !c.is_empty())
536 .map(str::to_string)
537 .collect()
538 })
539 .unwrap_or_default();
540
541 let mut single_int_pk = None;
544 if let Some(pk_col) = src.query_scalar(&pk_sql)? {
545 let type_sql = format!(
548 "SELECT t.name FROM sys.columns c \
549 JOIN sys.types t ON t.user_type_id = c.user_type_id \
550 JOIN sys.objects o ON o.object_id = c.object_id \
551 JOIN sys.schemas s ON s.schema_id = o.schema_id \
552 WHERE s.name = N'{}' AND o.name = N'{}' AND c.name = N'{}'",
553 schema.replace('\'', "''"),
554 table.replace('\'', "''"),
555 pk_col.replace('\'', "''")
556 );
557 if let Some(ty) = src.query_scalar(&type_sql)?
558 && matches!(ty.as_str(), "tinyint" | "smallint" | "int" | "bigint")
559 {
560 single_int_pk = Some(pk_col);
561 }
562 }
563
564 Ok(TableIntrospection {
565 single_int_pk,
566 keyset_keys,
567 row_estimate,
568 avg_row_bytes: None,
569 })
570}