fix(mssql): support cte in virtual tables (#18567)

* Fix for handling regular CTE queries with MSSQL,#8074

* Moved the get_cte_query function from mssql.py to base.py for using irrespetcive of dbengine

* Fix for handling regular CTE queries with MSSQL,#8074

* Moved the get_cte_query function from mssql.py to base.py for using irrespetcive of dbengine

* Unit test added for the db engine CTE SQL parsing.

Unit test added for the db engine CTE SQL parsing.  Removed additional spaces from the CTE parsing SQL generation.

* implement in sqla model

* lint + cleanup

Co-authored-by: Ville Brofeldt <ville.v.brofeldt@gmail.com>

tests/unit_tests/db_engine_specs/test_base.py

@@ -16,7 +16,11 @@
 # under the License.
 # pylint: disable=unused-argument, import-outside-toplevel, protected-access
 
+from textwrap import dedent
+
+import pytest
 from flask.ctx import AppContext
+from sqlalchemy.types import TypeEngine
 
 
 def test_get_text_clause_with_colon(app_context: AppContext) -> None:
@@ -56,3 +60,42 @@ def test_parse_sql_multi_statement(app_context: AppContext) -> None:
         "SELECT foo FROM tbl1",
         "SELECT bar FROM tbl2",
     ]
+
+
+@pytest.mark.parametrize(
+    "original,expected",
+    [
+        (
+            dedent(
+                """
+with currency as
+(
+select 'INR' as cur
+)
+select * from currency
+"""
+            ),
+            None,
+        ),
+        ("SELECT 1 as cnt", None,),
+        (
+            dedent(
+                """
+select 'INR' as cur
+union
+select 'AUD' as cur
+union
+select 'USD' as cur
+"""
+            ),
+            None,
+        ),
+    ],
+)
+def test_cte_query_parsing(
+    app_context: AppContext, original: TypeEngine, expected: str
+) -> None:
+    from superset.db_engine_specs.base import BaseEngineSpec
+
+    actual = BaseEngineSpec.get_cte_query(original)
+    assert actual == expected

tests/unit_tests/db_engine_specs/test_mssql.py

@@ -180,6 +180,57 @@ def test_column_datatype_to_string(
     assert actual == expected
 
 
+@pytest.mark.parametrize(
+    "original,expected",
+    [
+        (
+            dedent(
+                """
+with currency as (
+select 'INR' as cur
+),
+currency_2 as (
+select 'EUR' as cur
+)
+select * from currency union all select * from currency_2
+"""
+            ),
+            dedent(
+                """WITH currency as (
+select 'INR' as cur
+),
+currency_2 as (
+select 'EUR' as cur
+),
+__cte AS (
+select * from currency union all select * from currency_2
+)"""
+            ),
+        ),
+        ("SELECT 1 as cnt", None,),
+        (
+            dedent(
+                """
+select 'INR' as cur
+union
+select 'AUD' as cur
+union
+select 'USD' as cur
+"""
+            ),
+            None,
+        ),
+    ],
+)
+def test_cte_query_parsing(
+    app_context: AppContext, original: TypeEngine, expected: str
+) -> None:
+    from superset.db_engine_specs.mssql import MssqlEngineSpec
+
+    actual = MssqlEngineSpec.get_cte_query(original)
+    assert actual == expected
+
+
 def test_extract_errors(app_context: AppContext) -> None:
     """
     Test that custom error messages are extracted correctly.