Skip to main content

risingwave_connector/connector_common/
sql_server.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use anyhow::{Context, anyhow};
16use futures::TryStreamExt;
17use risingwave_common::types::DataType;
18use tiberius::{AuthMethod, Client, Config, EncryptionLevel};
19use tokio::net::TcpStream;
20use tokio_util::compat::TokioAsyncWriteCompatExt;
21
22pub type MssqlClient = Client<tokio_util::compat::Compat<TcpStream>>;
23
24/// Configuration for connecting to a SQL Server instance.
25///
26/// Mirrors `PgConnectionConfig` for PostgreSQL. Used by the `mssql_query`
27/// table-valued function and any future SQL Server sink/CDC source that needs
28/// a generic SQL Server client (independent of the Debezium-based CDC pipeline).
29#[derive(Debug, Clone)]
30pub struct MssqlConnectionConfig {
31    pub host: String,
32    pub port: u16,
33    pub user: String,
34    pub password: String,
35    pub database: String,
36    /// SQL Server "encrypt" connection string option: `true` means full TLS
37    /// encryption for all traffic; `false` means only the login handshake is
38    /// encrypted (the underlying `tiberius` [`EncryptionLevel::Off`] semantics).
39    pub encrypt: bool,
40    /// SQL Server `TrustServerCertificate` connection string option. When
41    /// `true`, the server certificate is accepted as-is without chain
42    /// validation. Implies `encrypt = true` since TLS is required for
43    /// certificate trust to apply.
44    pub trust_cert: bool,
45}
46
47/// Create an authenticated `tiberius::Client` connected to a SQL Server
48/// instance. Uses tokio's TCP stream wrapped with `TokioAsyncWriteCompatExt`
49/// to bridge between tokio's async I/O and tiberius's `futures-io` traits.
50///
51/// Encryption semantics:
52/// - `trust_cert = true` forces TLS on (matches SQL Server behavior where
53///   `TrustServerCertificate=true` is only meaningful with encryption).
54/// - Otherwise encryption follows the `encrypt` flag directly.
55pub async fn create_mssql_client(config: &MssqlConnectionConfig) -> anyhow::Result<MssqlClient> {
56    let mut tiberius_config = Config::new();
57    tiberius_config.host(&config.host);
58    tiberius_config.port(config.port);
59    tiberius_config.database(&config.database);
60    // trust_cert=true implies encrypt=true (TLS is required for cert trust).
61    tiberius_config.encryption(if config.trust_cert || config.encrypt {
62        EncryptionLevel::Required
63    } else {
64        EncryptionLevel::Off
65    });
66    if config.trust_cert {
67        tiberius_config.trust_cert();
68    }
69    tiberius_config.authentication(AuthMethod::sql_server(&config.user, &config.password));
70
71    create_mssql_client_with_config(tiberius_config)
72        .await
73        .with_context(|| {
74            format!(
75                "failed to connect to sql server {}:{} as user {}",
76                config.host, config.port, config.user
77            )
78        })
79}
80
81async fn connect_mssql_tcp(config: &Config, after_routing: bool) -> anyhow::Result<TcpStream> {
82    let address = config.get_addr();
83    let phase = if after_routing { " after routing" } else { "" };
84    let tcp = TcpStream::connect(&address)
85        .await
86        .with_context(|| format!("failed to connect to sql server at {address}{phase}"))?;
87    tcp.set_nodelay(true)
88        .with_context(|| format!("failed to set TCP_NODELAY for sql server at {address}{phase}"))?;
89    Ok(tcp)
90}
91
92/// Connect using a fully built Tiberius configuration.
93///
94/// SQL Server gateways, notably Azure SQL, can return one routing redirect
95/// during login. Follow that redirect here so query, CDC, and sink callers all
96/// use identical connection behavior.
97pub async fn create_mssql_client_with_config(mut config: Config) -> anyhow::Result<MssqlClient> {
98    let tcp = connect_mssql_tcp(&config, false).await?;
99    let connect_result = Client::connect(config.clone(), tcp.compat_write()).await;
100
101    match connect_result {
102        Ok(client) => Ok(client),
103        Err(tiberius::error::Error::Routing { host, port }) => {
104            config.host(&host);
105            config.port(port);
106            let tcp = connect_mssql_tcp(&config, true).await?;
107            Client::connect(config, tcp.compat_write())
108                .await
109                .context("failed to authenticate with sql server after routing")
110        }
111        Err(error) => Err(error).context("failed to authenticate with sql server"),
112    }
113}
114
115/// Discover the projected schema of an arbitrary user query against SQL Server.
116///
117/// Wraps the user query in `sp_describe_first_result_set`, which returns one
118/// row per result column with `name` and `system_type_name` (e.g. `int`,
119/// `nvarchar(50)`, `decimal(18,2)`). Each row is mapped to a RisingWave
120/// [`DataType`].
121///
122/// The user query is passed to `sp_describe_first_result_set` as a bound
123/// parameter. It is data for the stored procedure rather than part of the
124/// outer `EXEC` statement.
125pub async fn describe_mssql_query(
126    config: &MssqlConnectionConfig,
127    user_query: &str,
128) -> anyhow::Result<Vec<(String, DataType)>> {
129    let mut client = create_mssql_client(config).await?;
130
131    let stream = client
132        .query(
133            "EXEC sys.sp_describe_first_result_set @tsql = @P1",
134            &[&user_query],
135        )
136        .await?;
137    let mut row_stream = stream.into_row_stream();
138
139    let mut rw_types = vec![];
140    while let Some(row) = row_stream.try_next().await? {
141        // Column ordinal 0 -> `is_hidden` (bit). 0 = visible, 1 = hidden
142        // (a column hidden from the client, e.g. an unused join key).
143        // Column ordinal 2 -> `name` (sysname / nvarchar). Empty for
144        // hidden or computed-but-unnamed columns.
145        // Column ordinal 5 -> `system_type_name` (e.g. "int", "nvarchar(50)",
146        // "decimal(18,2)").
147        let is_hidden = row.try_get::<bool, _>(0)?.unwrap_or(false);
148        if is_hidden {
149            // Skip hidden columns: they are not visible in the result set.
150            continue;
151        }
152        let name: &str = row.try_get::<&str, _>(2)?.unwrap_or("");
153        let type_name: &str = row.try_get::<&str, _>(5)?.unwrap_or("");
154
155        let data_type = mssql_type_to_rw_type(type_name)
156            .with_context(|| format!("unsupported column type {:?}", type_name))?;
157        // Visible but unnamed columns (e.g. `SELECT COUNT(*) FROM t`) get a
158        // synthetic `column_N` name where N is the 1-based column ordinal.
159        // We keep the same default-name convention as the CDC source.
160        let column_name = if name.is_empty() {
161            row.try_get::<i32, _>(1)?
162                .map(|column_ordinal| format!("column_{column_ordinal}"))
163                .unwrap_or_else(|| format!("column_{}", rw_types.len() + 1))
164        } else {
165            name.to_owned()
166        };
167        rw_types.push((column_name, data_type));
168    }
169
170    if rw_types.is_empty() {
171        // `sp_describe_first_result_set` returns zero rows when the statement
172        // produces no result set, or when every column is hidden. Refuse to
173        // build a zero-column table function: downstream code (e.g. the
174        // executor's per-row decoder) would silently drop every row because
175        // there are no fields to decode into.
176        return Err(anyhow!(
177            "the query does not produce any result column: {:?}",
178            user_query
179        ));
180    }
181
182    Ok(rw_types)
183}
184
185/// Map a SQL Server `system_type_name` (as returned by
186/// `sp_describe_first_result_set`) to a RisingWave [`DataType`].
187///
188/// The name is the human-readable form including precision/length parameters,
189/// e.g. `int`, `nvarchar(50)`, `decimal(18,2)`, `datetime2(7)`. We strip the
190/// parenthesized suffix and dispatch on the base type. SQL Server CDC reuses
191/// this function with the bare type names from `INFORMATION_SCHEMA.COLUMNS`,
192/// keeping query schema discovery and CDC snapshots consistent.
193pub fn mssql_type_to_rw_type(type_name: &str) -> anyhow::Result<DataType> {
194    let lower = type_name.to_lowercase();
195    // Strip precision/scale/length suffix, e.g. "nvarchar(50)" -> "nvarchar".
196    let base = lower.split_once('(').map(|(b, _)| b).unwrap_or(&lower);
197    let dtype = match base.trim() {
198        "bit" => DataType::Boolean,
199        "binary" | "varbinary" => DataType::Bytea,
200        "tinyint" | "smallint" => DataType::Int16,
201        "int" | "integer" => DataType::Int32,
202        "bigint" => DataType::Int64,
203        "real" => DataType::Float32,
204        "float" => DataType::Float64,
205        "decimal" | "numeric" | "money" | "smallmoney" => DataType::Decimal,
206        "date" => DataType::Date,
207        "time" => DataType::Time,
208        "datetime" | "datetime2" | "smalldatetime" => DataType::Timestamp,
209        "datetimeoffset" => DataType::Timestamptz,
210        "char" | "nchar" | "varchar" | "nvarchar" | "text" | "ntext" | "xml"
211        | "uniqueidentifier" => DataType::Varchar,
212        unknown => {
213            return Err(anyhow!(
214                "unsupported sql server type: {:?} (from system_type_name {:?})",
215                unknown,
216                type_name
217            ));
218        }
219    };
220    Ok(dtype)
221}
222
223#[cfg(test)]
224mod tests {
225    use risingwave_common::types::DataType;
226
227    use super::mssql_type_to_rw_type;
228
229    /// Verify that the common base SQL Server types map to the expected
230    /// RisingWave `DataType`s: int/bigint/bit/float/date/datetime2/datetimeoffset/uniqueidentifier.
231    #[test]
232    fn mssql_type_to_rw_type_str_base_types() {
233        assert_eq!(mssql_type_to_rw_type("int").unwrap(), DataType::Int32);
234        assert_eq!(mssql_type_to_rw_type("bigint").unwrap(), DataType::Int64);
235        assert_eq!(mssql_type_to_rw_type("bit").unwrap(), DataType::Boolean);
236        assert_eq!(mssql_type_to_rw_type("float").unwrap(), DataType::Float64);
237        assert_eq!(mssql_type_to_rw_type("date").unwrap(), DataType::Date);
238        assert_eq!(
239            mssql_type_to_rw_type("datetime2").unwrap(),
240            DataType::Timestamp
241        );
242        assert_eq!(
243            mssql_type_to_rw_type("datetimeoffset").unwrap(),
244            DataType::Timestamptz
245        );
246        assert_eq!(
247            mssql_type_to_rw_type("uniqueidentifier").unwrap(),
248            DataType::Varchar
249        );
250    }
251
252    /// Verify that parameterized SQL Server types (`nvarchar(50)`,
253    /// `VARCHAR(255)`, `decimal(18, 2)`, `datetime2(7)`) are correctly
254    /// stripped of their parenthesized suffix and mapped.
255    #[test]
256    fn mssql_type_to_rw_type_str_parameterized() {
257        assert_eq!(
258            mssql_type_to_rw_type("nvarchar(50)").unwrap(),
259            DataType::Varchar
260        );
261        assert_eq!(
262            mssql_type_to_rw_type("VARCHAR(255)").unwrap(),
263            DataType::Varchar
264        );
265        assert_eq!(
266            mssql_type_to_rw_type("decimal(18, 2)").unwrap(),
267            DataType::Decimal
268        );
269        assert_eq!(
270            mssql_type_to_rw_type("datetime2(7)").unwrap(),
271            DataType::Timestamp
272        );
273    }
274
275    /// Verify that unknown SQL Server types (`geometry`, `hierarchyid`,
276    /// `sql_variant`) yield an error rather than silently mapping to
277    /// the wrong type.
278    #[test]
279    fn mssql_type_to_rw_type_str_unknown() {
280        assert!(mssql_type_to_rw_type("geometry").is_err());
281        assert!(mssql_type_to_rw_type("hierarchyid").is_err());
282        assert!(mssql_type_to_rw_type("sql_variant").is_err());
283    }
284}