risingwave_connector/connector_common/
sql_server.rs1use 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#[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 pub encrypt: bool,
40 pub trust_cert: bool,
45}
46
47pub 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 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
92pub 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
115pub 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 let is_hidden = row.try_get::<bool, _>(0)?.unwrap_or(false);
148 if is_hidden {
149 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 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 return Err(anyhow!(
177 "the query does not produce any result column: {:?}",
178 user_query
179 ));
180 }
181
182 Ok(rw_types)
183}
184
185pub fn mssql_type_to_rw_type(type_name: &str) -> anyhow::Result<DataType> {
194 let lower = type_name.to_lowercase();
195 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 #[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 #[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 #[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}