Fix!: preserve yaml order (#1054)

Author: tobymao

sqlmesh/core/schema_loader.py

@@ -10,7 +10,7 @@
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model import Model
 from sqlmesh.core.state_sync import StateReader
-from sqlmesh.utils.yaml import YAML
+from sqlmesh.utils import yaml
 
 logger = logging.getLogger(__name__)
 
@@ -62,4 +62,4 @@ def create_schema_file(
         ]
 
         with open(path, "w", encoding="utf-8") as file:
-            YAML().dump(schemas, file)
+            yaml.dump(schemas, file)

sqlmesh/dbt/builtin.py

@@ -204,7 +204,7 @@ def from_json(value: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any
 
 def to_yaml(value: t.Any, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
     try:
-        return yaml.dumps(value)
+        return yaml.dump(value)
     except (TypeError, YAMLError):
         return default
 

sqlmesh/dbt/common.py

@@ -34,7 +34,7 @@
 }
 
 
-def load_yaml(source: str | Path) -> t.OrderedDict:
+def load_yaml(source: str | Path) -> t.Dict:
     try:
         return load(source, render_jinja=False)
     except DuplicateKeyError as ex:

sqlmesh/dbt/profile.py

@@ -8,8 +8,8 @@
 from sqlmesh.dbt.common import PROJECT_FILENAME, load_yaml
 from sqlmesh.dbt.context import DbtContext
 from sqlmesh.dbt.target import TargetConfig
+from sqlmesh.utils import yaml
 from sqlmesh.utils.errors import ConfigError
-from sqlmesh.utils.yaml import dumps as dump_yaml
 
 logger = logging.getLogger(__name__)
 
@@ -105,7 +105,7 @@ def _read_profile(
                 f"Target '{target_name}' not specified in profiles for '{context.profile_name}'."
             )
 
-        target_fields = load_yaml(context.render(dump_yaml(outputs[target_name])))
+        target_fields = load_yaml(context.render(yaml.dump(outputs[target_name])))
         target = TargetConfig.load(
             {"name": target_name, "profile_name": context.profile_name, **target_fields}
         )

sqlmesh/magics.py

@@ -21,10 +21,8 @@
 from sqlmesh.core.dialect import format_model_expressions, parse
 from sqlmesh.core.model import load_model
 from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
+from sqlmesh.utils import yaml
 from sqlmesh.utils.errors import MagicError, MissingContextException
-from sqlmesh.utils.yaml import YAML
-from sqlmesh.utils.yaml import dumps as yaml_dumps
-from sqlmesh.utils.yaml import load as yaml_load
 
 CONTEXT_VARIABLE_NAMES = [
     "context",
@@ -208,8 +206,8 @@ def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
             return
 
         test = tests[model.name][args.test_name]
-        test_def = yaml_load(test_def_raw) if test_def_raw else test.body
-        test_def_output = yaml_dumps(test_def)
+        test_def = yaml.load(test_def_raw) if test_def_raw else test.body
+        test_def_output = yaml.dump(test_def)
 
         self._shell.set_next_input(
             "\n".join(
@@ -222,10 +220,10 @@ def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
         )
 
         with open(test.path, "r+", encoding="utf-8") as file:
-            content = yaml_load(file.read())
+            content = yaml.load(file.read())
             content[args.test_name] = test_def
             file.seek(0)
-            YAML().dump(content, file)
+            yaml.dump(content, file)
             file.truncate()
 
     @magic_arguments()

sqlmesh/utils/yaml.py

@@ -2,7 +2,6 @@
 
 import io
 import typing as t
-from collections import OrderedDict
 from os import getenv
 from pathlib import Path
 
@@ -23,7 +22,7 @@ def load(
     raise_if_empty: bool = True,
     render_jinja: bool = True,
     allow_duplicate_keys: bool = False,
-) -> t.OrderedDict:
+) -> t.Dict:
     """Loads a YAML object from either a raw string or a file."""
     path: t.Optional[Path] = None
 
@@ -42,13 +41,26 @@ def load(
         if raise_if_empty:
             error_path = f" '{path}'" if path else ""
             raise SQLMeshError(f"YAML source{error_path} can't be empty.")
-        return OrderedDict()
+        return {}
 
     return contents
 
 
-def dumps(value: dict | yaml.CommentedMap | OrderedDict) -> str:
-    """Dumps a ruamel.yaml loaded object and converts it into a string"""
+@t.overload
+def dump(value: t.Any, stream: io.IOBase) -> None:
+    ...
+
+
+@t.overload
+def dump(value: t.Any) -> str:
+    ...
+
+
+def dump(value: t.Any, stream: t.Optional[io.IOBase] = None) -> t.Optional[str]:
+    """Dumps a ruamel.yaml loaded object and converts it into a string or writes it to a stream."""
     result = io.StringIO()
-    yaml.YAML().dump(value, result)
+    yaml.YAML().dump(value, stream or result)
+
+    if stream:
+        return None
     return result.getvalue()

tests/core/test_schema_loader.py

@@ -95,8 +95,8 @@ def test_create_external_models(tmpdir, assert_exp_eq):
 def test_no_internal_model_conversion(tmp_path: Path, mocker: MockerFixture):
     engine_adapter_mock = mocker.Mock()
     engine_adapter_mock.columns.return_value = {
-        "a": exp.DataType.build("bigint"),
         "b": exp.DataType.build("text"),
+        "a": exp.DataType.build("bigint"),
     }
 
     state_reader_mock = mocker.Mock()
@@ -112,3 +112,4 @@ def test_no_internal_model_conversion(tmp_path: Path, mocker: MockerFixture):
 
     assert len(schema) == 1
     assert schema[0]["name"] == "tbl_c"
+    assert list(schema[0]["columns"]) == ["b", "a"]

tests/utils/test_yaml.py

@@ -17,7 +17,7 @@ def test_yaml() -> None:
       password: "{{ env_var('__SQLMESH_TEST_ENV_PASSWORD__') }}"
 """
 
-    assert contents == yaml.dumps(yaml.load(contents, render_jinja=False))
+    assert contents == yaml.dump(yaml.load(contents, render_jinja=False))
 
     expected_contents = """profile:
   target: prod
@@ -32,7 +32,7 @@ def test_yaml() -> None:
     os.environ["__SQLMESH_TEST_ENV_USER__"] = "user"
     os.environ["__SQLMESH_TEST_ENV_PASSWORD__"] = "password"
 
-    assert expected_contents == yaml.dumps(yaml.load(contents))
+    assert expected_contents == yaml.dump(yaml.load(contents))
 
     # Return the environment to its previous state
     del os.environ["__SQLMESH_TEST_ENV_USER__"]