[BUGFIX] Preserve quoting when serializing quoted table names (#11357)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: wookasz

great_expectations/datasource/fluent/sql_datasource.py

@@ -82,8 +82,6 @@
 )
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql import quoted_name  # noqa: TID251 # type-checking only
-
     # We re-import sqlalchemy here to make type-checking and our compatability layer
     # play nice with one another
     from great_expectations.compatibility import sqlalchemy
@@ -622,7 +620,7 @@ def get_batch(self, batch_request: BatchRequest) -> Batch:
         else:
             sql_partitioner = None
 
-        batch_spec_kwargs: dict[str, str | dict | None]
+        batch_spec_kwargs: Dict[str, str | dict | None]
         requests = self._fully_specified_batch_requests(batch_request)
         unsorted_metadata_dicts = [self._get_batch_metadata_from_batch_request(r) for r in requests]
 
@@ -921,7 +919,7 @@ def _validate_batch_request(self, batch_request: BatchRequest) -> None:
                 f"but actually has form:\n{pf(batch_request.dict())}\n"
             )
 
-    def _create_batch_spec_kwargs(self) -> dict[str, Any]:
+    def _create_batch_spec_kwargs(self) -> Dict[str, Any]:
         """Creates batch_spec_kwargs used to instantiate a SqlAlchemyDatasourceBatchSpec or RuntimeQueryBatchSpec
 
         This is called by get_batch to generate the batch.
@@ -974,7 +972,7 @@ def as_selectable(self) -> sqlalchemy.Selectable:
         return sa.select(sa.text(self.query.lstrip()[6:])).subquery()
 
     @override
-    def _create_batch_spec_kwargs(self) -> dict[str, Any]:
+    def _create_batch_spec_kwargs(self) -> Dict[str, Any]:
         return {
             "data_asset_name": self.name,
             "query": self.query,
@@ -1005,6 +1003,8 @@ class TableAsset(_SQLAsset):
     )
     schema_name: Optional[str] = None
 
+    _quote_character: Optional[str] = None
+
     @property
     def qualified_name(self) -> str:
         return f"{self.schema_name}.{self.table_name}" if self.schema_name else self.table_name
@@ -1019,9 +1019,7 @@ def _default_table_name(cls, table_name: str, values: dict, **kwargs) -> str:
         return validated_table_name
 
     @pydantic.validator("table_name")
-    def _resolve_quoted_name(cls, table_name: str) -> str | quoted_name:
-        table_name_is_quoted: bool = cls._is_bracketed_by_quotes(table_name)
-
+    def _resolve_quoted_name(cls, table_name: str, values: Dict[str, Any]) -> str:
         # We reimport sqlalchemy from our compatability layer because we make
         # quoted_name a top level import there.
         from great_expectations.compatibility import sqlalchemy
@@ -1030,19 +1028,34 @@ def _resolve_quoted_name(cls, table_name: str) -> str | quoted_name:
             if isinstance(table_name, sqlalchemy.quoted_name):
                 return table_name
 
-            if table_name_is_quoted:
+            quote: bool = cls._is_bracketed_by_quotes(table_name)
+
+            if quote:
                 # https://docs.sqlalchemy.org/en/20/core/sqlelement.html#sqlalchemy.sql.expression.quoted_name.quote
                 # Remove the quotes and add them back using the sqlalchemy.quoted_name function
                 # TODO: We need to handle nested quotes
-                table_name = table_name.strip("'").strip('"')
+                values["_quote_character"] = table_name[0]
+                quote = True
+                table_name = table_name.strip("".join(DEFAULT_QUOTE_CHARACTERS))
 
             return sqlalchemy.quoted_name(
                 value=table_name,
-                quote=table_name_is_quoted,
+                quote=quote,
             )
 
         return table_name
 
+    @override
+    def dict(self, **kwargs) -> Dict[str, Any]:
+        original_dict = super().dict(**kwargs)
+
+        # we need to ensure we retain the quotes when serializing quoted names
+        qc = self._quote_character
+        if qc is not None:
+            original_dict["table_name"] = f"{qc}{self.table_name}{qc}"
+
+        return original_dict
+
     @override
     def test_connection(self) -> None:
         """Test the connection for the TableAsset.
@@ -1081,7 +1094,7 @@ def as_selectable(self) -> sqlalchemy.Selectable:
         return sa.table(self.table_name, schema=self.schema_name)
 
     @override
-    def _create_batch_spec_kwargs(self) -> dict[str, Any]:
+    def _create_batch_spec_kwargs(self) -> Dict[str, Any]:
         return {
             "type": "table",
             "data_asset_name": self.name,
@@ -1091,7 +1104,7 @@ def _create_batch_spec_kwargs(self) -> dict[str, Any]:
         }
 
     @override
-    def _create_batch_spec(self, batch_spec_kwargs: dict) -> SqlAlchemyDatasourceBatchSpec:
+    def _create_batch_spec(self, batch_spec_kwargs: Dict) -> SqlAlchemyDatasourceBatchSpec:
         return SqlAlchemyDatasourceBatchSpec(**batch_spec_kwargs)
 
     @staticmethod
@@ -1135,7 +1148,7 @@ def _warn_for_more_specific_datasource_type(connection_string: str) -> None:
 
     connector: str = connection_string.split("://")[0].split("+")[0]
 
-    type_lookup_plus: dict[str, str] = {
+    type_lookup_plus: Dict[str, str] = {
         n: DataSourceManager.type_lookup[n].__name__
         for n in DataSourceManager.type_lookup.type_names()
     }

tests/datasource/fluent/test_sql_datasources.py

@@ -379,6 +379,23 @@ def test_unquoted_schema_names_are_added_as_lowercase(
         )
         assert table_asset.schema_name == schema_name.lower()
 
+    @pytest.mark.parametrize("table_name", ["my_table", "MY_TABLE", "My_Table"])
+    def test_unquoted_table_names_are_unquoted(
+        self,
+        sql_datasource_table_asset_test_connection_noop: SQLDatasource,
+        table_name: str,
+    ):
+        my_datasource: SQLDatasource = sql_datasource_table_asset_test_connection_noop
+
+        table_asset = my_datasource.add_table_asset(
+            name="my_table_asset",
+            table_name=table_name,
+            schema_name="my_schema",
+        )
+        assert isinstance(table_asset.table_name, sqlalchemy.quoted_name)
+        assert table_asset.table_name == table_name
+        assert not table_asset.table_name.quote
+
     @pytest.mark.parametrize(
         "schema_name",
         [
@@ -404,6 +421,55 @@ def test_quoted_schema_names_are_not_modified(
         )
         assert table_asset.schema_name == schema_name
 
+    @pytest.mark.parametrize(
+        "table_name",
+        [
+            '"my_table"',
+            '"MY_TABLE"',
+            '"My_Table"',
+            "'my_table'",
+            "'MY_TABLE'",
+            "'My_Table'",
+        ],
+    )
+    def test_quoted_table_names_are_quoted(
+        self,
+        sql_datasource_table_asset_test_connection_noop: SQLDatasource,
+        table_name: str,
+    ):
+        my_datasource: SQLDatasource = sql_datasource_table_asset_test_connection_noop
+
+        table_asset = my_datasource.add_table_asset(
+            name="my_table_asset",
+            table_name=table_name,
+            schema_name="my_schema",
+        )
+        assert isinstance(table_asset.table_name, sqlalchemy.quoted_name)
+        assert table_asset.table_name == table_name.lstrip("\"'").rstrip("\"'")
+        assert table_asset.table_name.quote
+
+    @pytest.mark.parametrize(
+        "table_name",
+        [
+            '"my_table"',
+            '"MY_TABLE"',
+            '"My_Table"',
+            "'my_table'",
+            "'MY_TABLE'",
+            "'My_Table'",
+            "my_table",
+            "MY_TABLE",
+            "My_Table",
+        ],
+    )
+    def test_table_name_serialization_preserves_quotes(
+        self,
+        table_name: str,
+    ):
+        table_asset = TableAsset(name="my_table_asset", table_name=table_name)
+        serialized = table_asset.dict()
+        assert serialized["table_name"] == table_name
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-vv"])