Skip to main content

connectorx/
sql.rs

1use crate::errors::ConnectorXError;
2#[cfg(feature = "src_oracle")]
3use crate::sources::oracle::OracleDialect;
4use fehler::{throw, throws};
5use log::{debug, trace, warn};
6use sqlparser::ast::helpers::attached_token::AttachedToken;
7use sqlparser::ast::{
8    BinaryOperator, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList,
9    FunctionArguments, GroupByExpr, Ident, LimitClause, ObjectName, Query, Select, SelectFlavor,
10    SelectItem, SetExpr, Statement, TableAlias, TableFactor, TableWithJoins, Value,
11    WildcardAdditionalOptions,
12};
13use sqlparser::dialect::Dialect;
14use sqlparser::parser::Parser;
15#[cfg(feature = "src_oracle")]
16use std::any::Any;
17
18#[derive(Debug, Clone)]
19pub enum CXQuery<Q = String> {
20    Naked(Q),   // The query directly comes from the user
21    Wrapped(Q), // The user query is already wrapped in a subquery
22}
23
24impl<Q: std::fmt::Display> std::fmt::Display for CXQuery<Q> {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        match self {
27            CXQuery::Naked(q) => write!(f, "{}", q),
28            CXQuery::Wrapped(q) => write!(f, "{}", q),
29        }
30    }
31}
32
33impl<Q: AsRef<str>> CXQuery<Q> {
34    pub fn as_str(&self) -> &str {
35        match self {
36            CXQuery::Naked(q) => q.as_ref(),
37            CXQuery::Wrapped(q) => q.as_ref(),
38        }
39    }
40}
41
42impl From<&str> for CXQuery {
43    fn from(s: &str) -> CXQuery<String> {
44        CXQuery::Naked(s.to_string())
45    }
46}
47
48impl From<&&str> for CXQuery {
49    fn from(s: &&str) -> CXQuery<String> {
50        CXQuery::Naked(s.to_string())
51    }
52}
53
54impl From<&String> for CXQuery {
55    fn from(s: &String) -> CXQuery {
56        CXQuery::Naked(s.clone())
57    }
58}
59
60impl From<&CXQuery> for CXQuery {
61    fn from(q: &CXQuery) -> CXQuery {
62        q.clone()
63    }
64}
65
66impl CXQuery<String> {
67    pub fn naked<Q: AsRef<str>>(q: Q) -> Self {
68        CXQuery::Naked(q.as_ref().to_string())
69    }
70}
71
72impl<Q: AsRef<str>> AsRef<str> for CXQuery<Q> {
73    fn as_ref(&self) -> &str {
74        match self {
75            CXQuery::Naked(q) => q.as_ref(),
76            CXQuery::Wrapped(q) => q.as_ref(),
77        }
78    }
79}
80
81impl<Q> CXQuery<Q> {
82    pub fn map<F, U>(&self, f: F) -> CXQuery<U>
83    where
84        F: Fn(&Q) -> U,
85    {
86        match self {
87            CXQuery::Naked(q) => CXQuery::Naked(f(q)),
88            CXQuery::Wrapped(q) => CXQuery::Wrapped(f(q)),
89        }
90    }
91}
92
93impl<Q, E> CXQuery<Result<Q, E>> {
94    pub fn result(self) -> Result<CXQuery<Q>, E> {
95        match self {
96            CXQuery::Naked(q) => q.map(CXQuery::Naked),
97            CXQuery::Wrapped(q) => q.map(CXQuery::Wrapped),
98        }
99    }
100}
101
102fn number(n: impl ToString) -> Expr {
103    Expr::value(Value::Number(n.to_string(), false))
104}
105
106fn function(name: &str, args: Vec<FunctionArg>) -> Expr {
107    Expr::Function(Function {
108        name: ObjectName::from(vec![Ident::new(name)]),
109        uses_odbc_syntax: false,
110        parameters: FunctionArguments::None,
111        args: FunctionArguments::List(FunctionArgumentList {
112            duplicate_treatment: None,
113            args,
114            clauses: vec![],
115        }),
116        filter: None,
117        null_treatment: None,
118        over: None,
119        within_group: vec![],
120    })
121}
122
123fn has_limit(query: &Query) -> bool {
124    match &query.limit_clause {
125        Some(LimitClause::LimitOffset { limit, .. }) => limit.is_some(),
126        Some(LimitClause::OffsetCommaLimit { .. }) => true,
127        None => false,
128    }
129}
130
131fn has_offset(query: &Query) -> bool {
132    match &query.limit_clause {
133        Some(LimitClause::LimitOffset { offset, .. }) => offset.is_some(),
134        Some(LimitClause::OffsetCommaLimit { .. }) => true,
135        None => false,
136    }
137}
138
139// set the LIMIT of a query, keeping any existing OFFSET
140fn set_limit(query: &mut Query, n: u64) {
141    match &mut query.limit_clause {
142        Some(LimitClause::LimitOffset { limit, .. }) => *limit = Some(number(n)),
143        Some(LimitClause::OffsetCommaLimit { limit, .. }) => *limit = number(n),
144        None => {
145            query.limit_clause = Some(LimitClause::LimitOffset {
146                limit: Some(number(n)),
147                offset: None,
148                limit_by: vec![],
149            })
150        }
151    }
152}
153
154// wrap a query into a derived table
155fn wrap_query(
156    query: &mut Query,
157    projection: Vec<SelectItem>,
158    selection: Option<Expr>,
159    tmp_tab_name: &str,
160) -> Statement {
161    let with = query.with.clone();
162    query.with = None;
163    let alias = if tmp_tab_name.is_empty() {
164        None
165    } else {
166        Some(TableAlias {
167            explicit: true,
168            name: Ident::new(tmp_tab_name),
169            columns: vec![],
170            at: None,
171        })
172    };
173    Statement::Query(Box::new(Query {
174        with,
175        locks: vec![],
176        body: Box::new(SetExpr::Select(Box::new(Select {
177            select_token: AttachedToken::empty(),
178            optimizer_hints: vec![],
179            distinct: None,
180            select_modifiers: None,
181            top: None,
182            top_before_distinct: false,
183            projection,
184            exclude: None,
185            into: None,
186            from: vec![TableWithJoins {
187                relation: TableFactor::Derived {
188                    lateral: false,
189                    subquery: Box::new(query.clone()),
190                    alias,
191                    sample: None,
192                },
193                joins: vec![],
194            }],
195            lateral_views: vec![],
196            prewhere: None,
197            selection,
198            connect_by: vec![],
199            group_by: GroupByExpr::Expressions(vec![], vec![]),
200            cluster_by: vec![],
201            distribute_by: vec![],
202            sort_by: vec![],
203            having: None,
204            named_window: vec![],
205            qualify: None,
206            window_before_qualify: false,
207            value_table_mode: None,
208            flavor: SelectFlavor::Standard,
209        }))),
210        order_by: None,
211        limit_clause: None,
212        fetch: None,
213        for_clause: None,
214        settings: None,
215        format_clause: None,
216        pipe_operators: vec![],
217    }))
218}
219
220trait StatementExt {
221    fn as_query(&self) -> Option<&Query>;
222}
223
224impl StatementExt for Statement {
225    fn as_query(&self) -> Option<&Query> {
226        match self {
227            Statement::Query(q) => Some(q),
228            _ => None,
229        }
230    }
231}
232
233trait QueryExt {
234    fn as_select_mut(&mut self) -> Option<&mut Select>;
235}
236
237impl QueryExt for Query {
238    fn as_select_mut(&mut self) -> Option<&mut Select> {
239        match *self.body {
240            SetExpr::Select(ref mut select) => Some(select),
241            _ => None,
242        }
243    }
244}
245
246#[throws(ConnectorXError)]
247pub fn count_query<T: Dialect>(sql: &CXQuery<String>, dialect: &T) -> CXQuery<String> {
248    trace!("Incoming query: {}", sql);
249
250    const COUNT_TMP_TAB_NAME: &str = "CXTMPTAB_COUNT";
251
252    #[allow(unused_mut)]
253    let mut table_alias = COUNT_TMP_TAB_NAME;
254
255    // HACK: Some dialect (e.g. Oracle) does not support "AS" for alias
256    #[cfg(feature = "src_oracle")]
257    if dialect.type_id() == (OracleDialect {}.type_id()) {
258        // table_alias = "";
259        return CXQuery::Wrapped(format!(
260            "SELECT COUNT(*) FROM ({}) {}",
261            sql.as_str(),
262            COUNT_TMP_TAB_NAME
263        ));
264    }
265
266    let tsql = match sql.map(|sql| Parser::parse_sql(dialect, sql)).result() {
267        Ok(ast) => {
268            let projection = vec![SelectItem::UnnamedExpr(function(
269                "count",
270                vec![FunctionArg::Unnamed(FunctionArgExpr::Wildcard)],
271            ))];
272            let ast_count: Statement = match ast {
273                CXQuery::Naked(ast) => {
274                    if ast.len() != 1 {
275                        throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
276                    }
277                    let mut query = ast[0]
278                        .as_query()
279                        .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
280                        .clone();
281                    if !has_offset(&query) {
282                        query.order_by = None; // mssql offset must appear with order by
283                    }
284                    let select = query
285                        .as_select_mut()
286                        .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?;
287                    select.sort_by = vec![];
288                    wrap_query(&mut query, projection, None, table_alias)
289                }
290                CXQuery::Wrapped(ast) => {
291                    if ast.len() != 1 {
292                        throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
293                    }
294                    let mut query = ast[0]
295                        .as_query()
296                        .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
297                        .clone();
298                    let select = query
299                        .as_select_mut()
300                        .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?;
301                    select.projection = projection;
302                    Statement::Query(Box::new(query))
303                }
304            };
305            format!("{}", ast_count)
306        }
307        Err(e) => {
308            warn!("parser error: {:?}, manually compose query string", e);
309            format!(
310                "SELECT COUNT(*) FROM ({}) as {}",
311                sql.as_str(),
312                COUNT_TMP_TAB_NAME
313            )
314        }
315    };
316
317    debug!("Transformed count query: {}", tsql);
318    CXQuery::Wrapped(tsql)
319}
320
321#[throws(ConnectorXError)]
322pub fn limit0_query<T: Dialect>(sql: &CXQuery<String>, dialect: &T) -> CXQuery<String> {
323    trace!("Incoming query: {}", sql);
324
325    let sql = match Parser::parse_sql(dialect, sql.as_str()) {
326        Ok(mut ast) => {
327            if ast.len() != 1 {
328                throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
329            }
330
331            match &mut ast[0] {
332                Statement::Query(q) => {
333                    set_limit(q, 0);
334                }
335                _ => throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string())),
336            };
337
338            format!("{}", ast[0])
339        }
340        Err(e) => {
341            warn!("parser error: {:?}, manually compose query string", e);
342            format!("{} LIMIT 0", sql.as_str())
343        }
344    };
345
346    debug!("Transformed limit 0 query: {}", sql);
347    CXQuery::Wrapped(sql)
348}
349
350/// Like limit0_query but with LIMIT 1. Used by SQLite where schema inference
351/// requires at least one row when decl_type is not available.
352#[throws(ConnectorXError)]
353pub fn limit1_query<T: Dialect>(sql: &CXQuery<String>, dialect: &T) -> CXQuery<String> {
354    trace!("Incoming query: {}", sql);
355
356    let sql = match Parser::parse_sql(dialect, sql.as_str()) {
357        Ok(mut ast) => {
358            if ast.len() != 1 {
359                throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
360            }
361
362            match &mut ast[0] {
363                Statement::Query(q) => {
364                    set_limit(q, 1);
365                }
366                _ => throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string())),
367            };
368
369            format!("{}", ast[0])
370        }
371        Err(e) => {
372            warn!("parser error: {:?}, manually compose query string", e);
373            format!("{} LIMIT 1", sql.as_str())
374        }
375    };
376
377    debug!("Transformed limit 1 query: {}", sql);
378    CXQuery::Wrapped(sql)
379}
380
381#[throws(ConnectorXError)]
382#[cfg(feature = "src_oracle")]
383pub fn limit0_query_oracle(sql: &CXQuery<String>) -> CXQuery<String> {
384    trace!("Incoming oracle query: {}", sql);
385
386    CXQuery::Wrapped(format!("SELECT * FROM ({}) WHERE 1=0", sql))
387
388    // let ast = Parser::parse_sql(&OracleDialect {}, sql.as_str())?;
389    // if ast.len() != 1 {
390    //     throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
391    // }
392    // let ast_part: Statement;
393    // let mut query = ast[0]
394    //     .as_query()
395    //     .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
396    //     .clone();
397
398    // let selection = Expr::BinaryOp {
399    //     left: Box::new(Expr::CompoundIdentifier(vec![Ident {
400    //         value: "rownum".to_string(),
401    //         quote_style: None,
402    //     }])),
403    //     op: BinaryOperator::Eq,
404    //     right: Box::new(Expr::Value(Value::Number("1".to_string(), false))),
405    // };
406    // ast_part = wrap_query(&mut query, vec![SelectItem::Wildcard], Some(selection), "");
407
408    // let tsql = format!("{}", ast_part);
409    // debug!("Transformed limit 1 query: {}", tsql);
410    // CXQuery::Wrapped(tsql)
411}
412
413#[throws(ConnectorXError)]
414pub fn single_col_partition_query<T: Dialect>(
415    sql: &str,
416    col: &str,
417    lower: i64,
418    upper: i64,
419    dialect: &T,
420) -> String {
421    trace!("Incoming query: {}", sql);
422    const PART_TMP_TAB_NAME: &str = "CXTMPTAB_PART";
423
424    #[allow(unused_mut)]
425    let mut table_alias = PART_TMP_TAB_NAME;
426    #[allow(unused_mut)]
427    let mut cid = Box::new(Expr::CompoundIdentifier(vec![
428        Ident::new(PART_TMP_TAB_NAME),
429        Ident::new(col),
430    ]));
431
432    // HACK: Some dialect (e.g. Oracle) does not support "AS" for alias
433    #[cfg(feature = "src_oracle")]
434    if dialect.type_id() == (OracleDialect {}.type_id()) {
435        return format!("SELECT * FROM ({}) CXTMPTAB_PART WHERE CXTMPTAB_PART.{} >= {} AND CXTMPTAB_PART.{} < {}", sql, col, lower, col, upper);
436        // table_alias = "";
437        // cid = Box::new(Expr::Identifier(Ident {
438        //     value: col.to_string(),
439        //     quote_style: None,
440        // }));
441    }
442
443    let tsql = match Parser::parse_sql(dialect, sql) {
444        Ok(ast) => {
445            if ast.len() != 1 {
446                throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
447            }
448
449            let mut query = ast[0]
450                .as_query()
451                .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
452                .clone();
453
454            let select = query
455                .as_select_mut()
456                .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
457                .clone();
458
459            let lb = Expr::BinaryOp {
460                left: Box::new(number(lower)),
461                op: BinaryOperator::LtEq,
462                right: cid.clone(),
463            };
464
465            let ub = Expr::BinaryOp {
466                left: cid,
467                op: BinaryOperator::Lt,
468                right: Box::new(number(upper)),
469            };
470
471            let selection = Expr::BinaryOp {
472                left: Box::new(lb),
473                op: BinaryOperator::And,
474                right: Box::new(ub),
475            };
476
477            if !has_limit(&query) && select.top.is_none() && query.order_by.is_some() {
478                // order by in a partition query does not make sense because partition is unordered.
479                // clear the order by beceause mssql does not support order by in a derived table.
480                // also order by in the derived table does not make any difference.
481                query.order_by = None;
482            }
483
484            let ast_part: Statement = wrap_query(
485                &mut query,
486                vec![SelectItem::Wildcard(WildcardAdditionalOptions::default())],
487                Some(selection),
488                table_alias,
489            );
490            format!("{}", ast_part)
491        }
492        Err(e) => {
493            warn!("parser error: {:?}, manually compose query string", e);
494            format!("SELECT * FROM ({}) AS CXTMPTAB_PART WHERE CXTMPTAB_PART.{} >= {} AND CXTMPTAB_PART.{} < {}", sql, col, lower, col, upper)
495        }
496    };
497
498    debug!("Transformed single column partition query: {}", tsql);
499    tsql
500}
501
502#[throws(ConnectorXError)]
503pub fn get_partition_range_query<T: Dialect>(sql: &str, col: &str, dialect: &T) -> String {
504    trace!("Incoming query: {}", sql);
505    const RANGE_TMP_TAB_NAME: &str = "CXTMPTAB_RANGE";
506
507    #[allow(unused_mut)]
508    let mut table_alias = RANGE_TMP_TAB_NAME;
509    #[allow(unused_mut)]
510    let mut args = vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
511        Expr::CompoundIdentifier(vec![Ident::new(RANGE_TMP_TAB_NAME), Ident::new(col)]),
512    ))];
513
514    // HACK: Some dialect (e.g. Oracle) does not support "AS" for alias
515    #[cfg(feature = "src_oracle")]
516    if dialect.type_id() == (OracleDialect {}.type_id()) {
517        return format!(
518            "SELECT MIN({}.{}) as min, MAX({}.{}) as max FROM ({}) {}",
519            RANGE_TMP_TAB_NAME, col, RANGE_TMP_TAB_NAME, col, sql, RANGE_TMP_TAB_NAME
520        );
521        // table_alias = "";
522        // args = vec![FunctionArg::Unnamed(Expr::Identifier(Ident {
523        //     value: col.to_string(),
524        //     quote_style: None,
525        // }))];
526    }
527
528    let tsql = match Parser::parse_sql(dialect, sql) {
529        Ok(ast) => {
530            if ast.len() != 1 {
531                throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
532            }
533
534            let mut query = ast[0]
535                .as_query()
536                .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
537                .clone();
538
539            if !has_limit(&query) && !has_offset(&query) {
540                query.order_by = None; // only omit orderby when there is no limit and offset in the query
541            }
542            let projection = vec![
543                SelectItem::UnnamedExpr(function("min", args.clone())),
544                SelectItem::UnnamedExpr(function("max", args)),
545            ];
546            let ast_range: Statement = wrap_query(&mut query, projection, None, table_alias);
547            format!("{}", ast_range)
548        }
549        Err(e) => {
550            warn!("parser error: {:?}, manually compose query string", e);
551            format!(
552                "SELECT MIN({}.{}) as min, MAX({}.{}) as max FROM ({}) AS {}",
553                RANGE_TMP_TAB_NAME, col, RANGE_TMP_TAB_NAME, col, sql, RANGE_TMP_TAB_NAME
554            )
555        }
556    };
557
558    debug!("Transformed partition range query: {}", tsql);
559    tsql
560}
561
562#[throws(ConnectorXError)]
563pub fn get_partition_range_query_sep<T: Dialect>(
564    sql: &str,
565    col: &str,
566    dialect: &T,
567) -> (String, String) {
568    trace!("Incoming query: {}", sql);
569    const RANGE_TMP_TAB_NAME: &str = "CXTMPTAB_RANGE";
570
571    let (sql_min, sql_max) = match Parser::parse_sql(dialect, sql) {
572        Ok(ast) => {
573            if ast.len() != 1 {
574                throw!(ConnectorXError::SqlQueryNotSupported(sql.to_string()));
575            }
576
577            let mut query = ast[0]
578                .as_query()
579                .ok_or_else(|| ConnectorXError::SqlQueryNotSupported(sql.to_string()))?
580                .clone();
581
582            query.order_by = None;
583            let args = vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
584                Expr::CompoundIdentifier(vec![Ident::new(RANGE_TMP_TAB_NAME), Ident::new(col)]),
585            ))];
586            let min_proj = vec![SelectItem::UnnamedExpr(function("min", args.clone()))];
587            let max_proj = vec![SelectItem::UnnamedExpr(function("max", args))];
588            let ast_range_min: Statement =
589                wrap_query(&mut query.clone(), min_proj, None, RANGE_TMP_TAB_NAME);
590            let ast_range_max: Statement =
591                wrap_query(&mut query, max_proj, None, RANGE_TMP_TAB_NAME);
592            (format!("{}", ast_range_min), format!("{}", ast_range_max))
593        }
594        Err(e) => {
595            warn!("parser error: {:?}, manually compose query string", e);
596            (
597                format!(
598                    "SELECT MIN({}.{}) as min FROM ({}) AS {}",
599                    RANGE_TMP_TAB_NAME, col, sql, RANGE_TMP_TAB_NAME
600                ),
601                format!(
602                    "SELECT MAX({}.{}) as max FROM ({}) AS {}",
603                    RANGE_TMP_TAB_NAME, col, sql, RANGE_TMP_TAB_NAME
604                ),
605            )
606        }
607    };
608    debug!(
609        "Transformed separated partition range query: {}, {}",
610        sql_min, sql_max
611    );
612    (sql_min, sql_max)
613}
614
615#[cfg(test)]
616mod tests {
617    use super::{
618        count_query, get_partition_range_query, get_partition_range_query_sep, limit0_query,
619        limit1_query, single_col_partition_query, CXQuery,
620    };
621    use crate::errors::ConnectorXError;
622    use sqlparser::dialect::{GenericDialect, MsSqlDialect, MySqlDialect};
623
624    #[test]
625    fn preserves_query_variant_and_maps_values() {
626        let naked = CXQuery::naked("select 1");
627        assert_eq!(naked.as_str(), "select 1");
628        assert!(matches!(naked.map(|query| query.len()), CXQuery::Naked(8)));
629        assert_eq!(
630            CXQuery::Wrapped("select 1".to_string()).to_string(),
631            "select 1"
632        );
633    }
634
635    #[test]
636    fn unwraps_successful_and_failed_query_results() {
637        let naked: CXQuery<Result<&str, &str>> = CXQuery::Naked(Ok("select 1"));
638        assert_eq!(naked.result().unwrap().as_str(), "select 1");
639
640        let wrapped: CXQuery<Result<&str, &str>> = CXQuery::Wrapped(Err("bad"));
641        assert!(matches!(wrapped.result(), Err("bad")));
642    }
643
644    #[test]
645    fn transforms_count_and_limit_queries() {
646        let query = CXQuery::naked("SELECT id FROM items ORDER BY id");
647        let dialect = GenericDialect {};
648        assert!(count_query(&query, &dialect)
649            .unwrap()
650            .as_str()
651            .to_ascii_lowercase()
652            .contains("count"));
653        assert!(limit0_query(&query, &dialect)
654            .unwrap()
655            .as_str()
656            .contains("LIMIT 0"));
657        assert!(limit1_query(&query, &dialect)
658            .unwrap()
659            .as_str()
660            .contains("LIMIT 1"));
661    }
662
663    #[test]
664    fn uses_fallback_for_unparseable_limit_query() {
665        let query = CXQuery::naked("SELECT FROM");
666        assert_eq!(
667            limit0_query(&query, &GenericDialect {}).unwrap().as_str(),
668            "SELECT FROM LIMIT 0"
669        );
670        assert_eq!(
671            limit1_query(&query, &GenericDialect {}).unwrap().as_str(),
672            "SELECT FROM LIMIT 1"
673        );
674    }
675
676    #[test]
677    fn rejects_multiple_and_non_select_queries() {
678        let dialect = GenericDialect {};
679        assert!(matches!(
680            count_query(&CXQuery::naked("SELECT 1; SELECT 2"), &dialect),
681            Err(ConnectorXError::SqlQueryNotSupported(_))
682        ));
683        assert!(matches!(
684            limit0_query(&CXQuery::naked("INSERT INTO items VALUES (1)"), &dialect),
685            Err(ConnectorXError::SqlQueryNotSupported(_))
686        ));
687    }
688
689    #[test]
690    fn creates_partition_and_range_queries() {
691        let dialect = GenericDialect {};
692        let partition = single_col_partition_query(
693            "SELECT id, value FROM items ORDER BY id",
694            "id",
695            10,
696            20,
697            &dialect,
698        )
699        .unwrap();
700        assert!(partition.contains("CXTMPTAB_PART"));
701        assert!(partition.contains("10"));
702        assert!(partition.contains("20"));
703        assert!(!partition.contains("ORDER BY id"));
704
705        let range =
706            get_partition_range_query("SELECT id FROM items ORDER BY id", "id", &dialect).unwrap();
707        let range_lower = range.to_ascii_lowercase();
708        assert!(range_lower.contains("min"));
709        assert!(range_lower.contains("max"));
710        assert!(range.contains("CXTMPTAB_RANGE"));
711
712        let (min, max) =
713            get_partition_range_query_sep("SELECT id FROM items ORDER BY id", "id", &dialect)
714                .unwrap();
715        assert!(min.to_ascii_lowercase().contains("min"));
716        assert!(min.contains("CXTMPTAB_RANGE"));
717        assert!(max.to_ascii_lowercase().contains("max"));
718        assert!(max.contains("CXTMPTAB_RANGE"));
719    }
720
721    #[test]
722    fn clears_mssql_ordering_only_when_safe() {
723        let query = single_col_partition_query(
724            "SELECT id FROM items ORDER BY id",
725            "id",
726            0,
727            1,
728            &MsSqlDialect {},
729        )
730        .unwrap();
731        assert!(!query.contains("ORDER BY id"));
732    }
733
734    #[test]
735    fn limit_queries_keep_existing_offset() {
736        let query = CXQuery::naked("SELECT id FROM items ORDER BY id LIMIT 5 OFFSET 2");
737        assert_eq!(
738            limit0_query(&query, &GenericDialect {}).unwrap().as_str(),
739            "SELECT id FROM items ORDER BY id LIMIT 0 OFFSET 2"
740        );
741        let query = CXQuery::naked("SELECT id FROM items LIMIT 2, 10");
742        assert_eq!(
743            limit1_query(&query, &MySqlDialect {}).unwrap().as_str(),
744            "SELECT id FROM items LIMIT 2, 1"
745        );
746    }
747
748    #[test]
749    fn keeps_ordering_when_query_has_limit_or_offset() {
750        let sql = "SELECT id FROM items ORDER BY id LIMIT 5 OFFSET 2";
751        let dialect = GenericDialect {};
752        assert!(count_query(&CXQuery::naked(sql), &dialect)
753            .unwrap()
754            .as_str()
755            .contains("ORDER BY id"));
756        assert!(single_col_partition_query(sql, "id", 0, 1, &dialect)
757            .unwrap()
758            .contains("ORDER BY id"));
759        assert!(get_partition_range_query(sql, "id", &dialect)
760            .unwrap()
761            .contains("ORDER BY id"));
762    }
763}