[FEATURE] Support asymmetric quoted identifiers in dialect quoting (#11652)

ashmortar · pre-commit-ci[bot] · web-flow · commit 1b24781511ee · 2026-02-13T18:34:12.000Z
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/great_expectations/execution_engine/sqlalchemy_dialect.py b/great_expectations/execution_engine/sqlalchemy_dialect.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Final, List, Literal, Mapping, Union, overload
+from typing import Any, Final, List, Mapping, Union, overload
 
 from great_expectations.compatibility.sqlalchemy import quoted_name
 from great_expectations.compatibility.typing_extensions import override
@@ -63,30 +63,33 @@ def get_all_dialects(cls) -> List[GXSqlDialect]:
         return [dialect for dialect in cls if dialect != GXSqlDialect.OTHER]
 
 
-DIALECT_IDENTIFIER_QUOTE_STRINGS: Final[Mapping[GXSqlDialect, Literal['"', "`"]]] = {
-    # TODO: add other dialects
-    GXSqlDialect.DATABRICKS: "`",
-    GXSqlDialect.MYSQL: "`",
-    GXSqlDialect.POSTGRESQL: '"',
-    GXSqlDialect.SNOWFLAKE: '"',
-    GXSqlDialect.SQLITE: '"',
-    GXSqlDialect.TRINO: "`",
+DIALECT_IDENTIFIER_QUOTE_STRINGS: Final[Mapping[GXSqlDialect, tuple[str, str]]] = {
+    GXSqlDialect.DATABRICKS: ("`", "`"),
+    GXSqlDialect.MSSQL: ("[", "]"),
+    GXSqlDialect.MYSQL: ("`", "`"),
+    GXSqlDialect.POSTGRESQL: ('"', '"'),
+    GXSqlDialect.SNOWFLAKE: ('"', '"'),
+    GXSqlDialect.SQLITE: ('"', '"'),
+    GXSqlDialect.TRINO: ("`", "`"),
 }
 
 
 def quote_str(unquoted_identifier: str, dialect: GXSqlDialect) -> str:
     """Quote a string using the specified dialect's quote character."""
-    quote_char = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect]
-    if unquoted_identifier.startswith(quote_char) or unquoted_identifier.endswith(quote_char):
+    open_q, close_q = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect]
+    if unquoted_identifier.startswith(open_q) or unquoted_identifier.endswith(close_q):
         raise ValueError(  # noqa: TRY003 # FIXME CoP
-            f"Identifier {unquoted_identifier} already uses quote character {quote_char}"
+            f"Identifier {unquoted_identifier} already uses quote characters {open_q}{close_q}"
         )
-    return f"{quote_char}{unquoted_identifier}{quote_char}"
+    return f"{open_q}{unquoted_identifier}{close_q}"
 
 
 def _strip_quotes(s: str, dialect: GXSqlDialect) -> str:
-    quote_str = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect]
-    if s.startswith(quote_str) and s.endswith(quote_str):
+    open_q, close_q = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect]
+    if s.startswith(open_q) and s.endswith(close_q):
+        return s[len(open_q) : -len(close_q)]
+    # MSSQL also accepts double quotes when QUOTED_IDENTIFIER is ON
+    if dialect == GXSqlDialect.MSSQL and s.startswith('"') and s.endswith('"'):
         return s[1:-1]
     return s
 
diff --git a/tests/execution_engine/test_sqlalchemy_dialect.py b/tests/execution_engine/test_sqlalchemy_dialect.py
@@ -1,6 +1,11 @@
 import pytest
 
-from great_expectations.execution_engine.sqlalchemy_dialect import GXSqlDialect
+from great_expectations.execution_engine.sqlalchemy_dialect import (
+    GXSqlDialect,
+    _strip_quotes,
+    quote_str,
+    wrap_identifier,
+)
 
 
 @pytest.mark.unit
@@ -31,3 +36,93 @@ def test_get_all_dialect_names_no_other_dialects():
 @pytest.mark.unit
 def test_get_all_dialects_no_other_dialects():
     assert GXSqlDialect.OTHER not in GXSqlDialect.get_all_dialects()
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dialect,expected",
+    [
+        (GXSqlDialect.DATABRICKS, "`col`"),
+        (GXSqlDialect.MSSQL, "[col]"),
+        (GXSqlDialect.MYSQL, "`col`"),
+        (GXSqlDialect.POSTGRESQL, '"col"'),
+        (GXSqlDialect.SNOWFLAKE, '"col"'),
+        (GXSqlDialect.SQLITE, '"col"'),
+        (GXSqlDialect.TRINO, "`col`"),
+    ],
+)
+def test_quote_str(dialect, expected):
+    assert quote_str("col", dialect) == expected
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dialect,quoted_input",
+    [
+        (GXSqlDialect.DATABRICKS, "`col`"),
+        (GXSqlDialect.MSSQL, "[col]"),
+        (GXSqlDialect.MYSQL, "`col`"),
+        (GXSqlDialect.POSTGRESQL, '"col"'),
+        (GXSqlDialect.SNOWFLAKE, '"col"'),
+        (GXSqlDialect.SQLITE, '"col"'),
+        (GXSqlDialect.TRINO, "`col`"),
+    ],
+)
+def test_quote_str_already_quoted_raises(dialect, quoted_input):
+    with pytest.raises(ValueError, match="already uses quote characters"):
+        quote_str(quoted_input, dialect)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dialect,quoted_input",
+    [
+        (GXSqlDialect.DATABRICKS, "`col`"),
+        (GXSqlDialect.MSSQL, "[col]"),
+        (GXSqlDialect.MSSQL, '"col"'),
+        (GXSqlDialect.MYSQL, "`col`"),
+        (GXSqlDialect.POSTGRESQL, '"col"'),
+        (GXSqlDialect.SNOWFLAKE, '"col"'),
+        (GXSqlDialect.SQLITE, '"col"'),
+        (GXSqlDialect.TRINO, "`col`"),
+    ],
+)
+def test_strip_quotes(dialect, quoted_input):
+    assert _strip_quotes(quoted_input, dialect) == "col"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dialect",
+    [
+        GXSqlDialect.DATABRICKS,
+        GXSqlDialect.MSSQL,
+        GXSqlDialect.MYSQL,
+        GXSqlDialect.POSTGRESQL,
+        GXSqlDialect.SNOWFLAKE,
+        GXSqlDialect.SQLITE,
+        GXSqlDialect.TRINO,
+    ],
+)
+def test_strip_quotes_unquoted_noop(dialect):
+    assert _strip_quotes("col", dialect) == "col"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dialect,quoted_input",
+    [
+        (GXSqlDialect.DATABRICKS, "`col`"),
+        (GXSqlDialect.MSSQL, "[col]"),
+        (GXSqlDialect.MSSQL, '"col"'),
+        (GXSqlDialect.MYSQL, "`col`"),
+        (GXSqlDialect.POSTGRESQL, '"col"'),
+        (GXSqlDialect.SNOWFLAKE, '"col"'),
+        (GXSqlDialect.SQLITE, '"col"'),
+        (GXSqlDialect.TRINO, "`col`"),
+    ],
+)
+def test_wrap_identifier_strips_quotes(dialect, quoted_input):
+    result = wrap_identifier(quoted_input, dialect=dialect)
+    assert str(result) == "col"
+    assert result.quote is True
diff --git a/tests/integration/fluent/test_sql_datasources.py b/tests/integration/fluent/test_sql_datasources.py
@@ -45,6 +45,9 @@
 )
 from great_expectations.core.expectation_suite import ExpectationSuite
 from great_expectations.core.validation_definition import ValidationDefinition
+from great_expectations.datasource.fluent.sql_server_datasource import (
+    SQLServerAuthConnectionDetails,
+)
 from great_expectations.execution_engine.sqlalchemy_dialect import (
     DIALECT_IDENTIFIER_QUOTE_STRINGS,
     GXSqlDialect,
@@ -64,6 +67,9 @@
         SQLDatasource,
         SqliteDatasource,
     )
+    from great_expectations.datasource.fluent.sql_server_datasource import (
+        SQLServerDatasource,
+    )
     from great_expectations.execution_engine import SqlAlchemyExecutionEngine
 
 TERMINAL_WIDTH: Final = shutil.get_terminal_size().columns
@@ -83,7 +89,7 @@
 # sqlite db files should be using fresh tmp_path on every test
 DO_NOT_DROP_TABLES: set[str] = {"sqlite"}
 
-DatabaseType: TypeAlias = Literal["postgres", "sqlite", "trino"]
+DatabaseType: TypeAlias = Literal["postgres", "sqlite", "trino", "mssql"]
 TableNameCase: TypeAlias = Literal[
     "quoted_lower",
     "quoted_mixed",
@@ -122,6 +128,14 @@
         "quoted_mixed": f'"{TEST_TABLE_NAME.title()}"',
         "unquoted_mixed": TEST_TABLE_NAME.title(),
     },
+    "mssql": {
+        "unquoted_lower": TEST_TABLE_NAME.lower(),
+        "quoted_lower": f"[{TEST_TABLE_NAME.lower()}]",
+        "unquoted_upper": TEST_TABLE_NAME.upper(),
+        "quoted_upper": f"[{TEST_TABLE_NAME.upper()}]",
+        "quoted_mixed": f"[{TEST_TABLE_NAME.title()}]",
+        "unquoted_mixed": TEST_TABLE_NAME.title(),
+    },
 }
 
 # column names
@@ -343,6 +357,24 @@ def __call__(
     ) -> None: ...
 
 
+def _create_schema_ddl(schema: str, is_mssql: bool) -> str:
+    if is_mssql:
+        return (
+            f"IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{schema}')"
+            f" EXEC('CREATE SCHEMA {schema}')"
+        )
+    return f"CREATE SCHEMA IF NOT EXISTS {schema}"
+
+
+def _create_table_ddl(qualified_table_name: str, table_columns: str, is_mssql: bool) -> str:
+    if is_mssql:
+        return (
+            f"IF OBJECT_ID(N'{qualified_table_name}', N'U') IS NULL"
+            f" CREATE TABLE {qualified_table_name}{table_columns}"
+        )
+    return f"CREATE TABLE IF NOT EXISTS {qualified_table_name}{table_columns}"
+
+
 @pytest.fixture(
     scope="class",
 )
@@ -382,6 +414,7 @@ def _table_factory(
         )
         dialect = GXSqlDialect(sa_engine.dialect.name)
         created_tables: list[dict[Literal["table_name", "schema"], str | None]] = []
+        is_mssql = dialect == GXSqlDialect.MSSQL
 
         with gx_engine.get_connection() as conn:
             quoted_upper_col: str = quote_str(QUOTED_UPPER_COL, dialect=dialect)
@@ -390,18 +423,18 @@ def _table_factory(
             quoted_mixed_case: str = quote_str(QUOTED_MIXED_CASE, dialect=dialect)
 
             if schema:
-                conn.execute(TextClause(f"CREATE SCHEMA IF NOT EXISTS {schema}"))
+                conn.execute(TextClause(_create_schema_ddl(schema, is_mssql)))
             for name in table_names:
                 qualified_table_name = f"{schema}.{name}" if schema else name
-                # TODO: use dialect specific quotes
-                create_tables: str = (
-                    f"CREATE TABLE IF NOT EXISTS {qualified_table_name}"
+                table_columns: str = (
                     " (id INTEGER, name VARCHAR(255),"
                     f" {quoted_upper_col} VARCHAR(255), {quoted_lower_col} VARCHAR(255),"
                     f" {UNQUOTED_UPPER_COL} VARCHAR(255), {UNQUOTED_LOWER_COL} VARCHAR(255),"
                     f" {quoted_mixed_case} VARCHAR(255), {quoted_w_dots} VARCHAR(255))"
                 )
-                conn.execute(TextClause(create_tables))
+                conn.execute(
+                    TextClause(_create_table_ddl(qualified_table_name, table_columns, is_mssql))
+                )
                 if data:
                     insert_data = (
                         f"INSERT INTO {qualified_table_name}"
@@ -470,6 +503,24 @@ def sqlite_ds(context: EphemeralDataContext, tmp_path: pathlib.Path) -> SqliteDa
     return ds
 
 
+@pytest.fixture
+def mssql_ds(context: EphemeralDataContext) -> SQLServerDatasource:
+    ds = context.data_sources.add_sql_server(
+        "mssql",
+        connection_string=SQLServerAuthConnectionDetails(
+            host="127.0.0.1",
+            port=1433,
+            database="test_ci",
+            schema="dbo",
+            username="sa",
+            password="ReallyStrongPwd1234%^&*",
+            driver="ODBC Driver 18 for SQL Server",
+            encrypt="Optional",
+        ),
+    )
+    return ds
+
+
 @pytest.fixture(
     params=[
         param(
@@ -551,6 +602,24 @@ def test_sqlite(
 
         sqlite_ds.add_table_asset(asset_name, table_name=table_name)
 
+    @pytest.mark.mssql
+    def test_mssql(
+        self,
+        mssql_ds: SQLServerDatasource,
+        asset_name: TableNameCase,
+        table_factory: TableFactory,
+    ):
+        table_name = TABLE_NAME_MAPPING["mssql"].get(asset_name)
+        if not table_name:
+            pytest.skip(f"no '{asset_name}' table_name for mssql")
+        # create table
+        table_factory(gx_engine=mssql_ds.get_execution_engine(), table_names={table_name})
+
+        table_names: list[str] = inspect(mssql_ds.get_engine()).get_table_names()
+        print(f"mssql tables:\n{pf(table_names)}))")
+
+        mssql_ds.add_table_asset(asset_name, table_name=table_name)
+
     @pytest.mark.filterwarnings(  # snowflake `add_table_asset` raises warning on passing a schema
         "once::great_expectations.datasource.fluent.GxDatasourceWarning"
     )
@@ -560,6 +629,7 @@ def test_sqlite(
             param("trino", None, marks=[pytest.mark.trino]),
             param("postgres", None, marks=[pytest.mark.postgresql]),
             param("sqlite", None, marks=[pytest.mark.sqlite]),
+            param("mssql", None, marks=[pytest.mark.mssql]),
         ],
     )
     def test_checkpoint_run(
@@ -625,10 +695,10 @@ def _is_quote_char_dialect_mismatch(
     dialect: GXSqlDialect,
     column_name: str | quoted_name,
 ) -> bool:
-    quote_char = column_name[0] if column_name[0] in ("'", '"', "`") else None
+    quote_char = column_name[0] if column_name[0] in ("'", '"', "`", "[") else None
     if quote_char:
-        dialect_quote_char = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect]
-        if quote_char != dialect_quote_char:
+        expected = DIALECT_IDENTIFIER_QUOTE_STRINGS[dialect][0]
+        if quote_char != expected:
             return True
     return False
 
@@ -639,8 +709,12 @@ def _raw_query_check_column_exists(
     gx_execution_engine: SqlAlchemyExecutionEngine,
 ) -> bool:
     """Use a simple 'SELECT {column_name_param} from {qualified_table_name};' query to check if the column exists.'"""  # noqa: E501 # FIXME CoP
-    with gx_execution_engine.get_connection() as connection:
+    dialect_name = gx_execution_engine.engine.dialect.name
+    if dialect_name == "mssql":
+        query = f"""SELECT TOP 1 {column_name_param} FROM {qualified_table_name};"""
+    else:
         query = f"""SELECT {column_name_param} FROM {qualified_table_name} LIMIT 1;"""
+    with gx_execution_engine.get_connection() as connection:
         print(f"query:\n  {query}")
         # an exception will be raised if the column does not exist
         try:
