fix: Register legacy DiscoveryFile object in Ruamel YAML instance (m…

…eltano#6543)

* fix yaml calls to old lib impl

* Register DiscoveryFile globally

* Make linter happy

Co-authored-by: Edgar R. M <edgarrm358@gmail.com>

src/meltano/core/discovery_file.py

@@ -0,0 +1,44 @@
+"""Legacy discovery file object."""
+
+from __future__ import annotations
+
+from .behavior.canonical import Canonical
+from .plugin import PluginDefinition, PluginType
+
+
+class DiscoveryFile(Canonical):
+    """A discovery file object."""
+
+    def __init__(self, version=1, **plugins):
+        """Create a new DiscoveryFile.
+
+        Args:
+            version: The version of the discovery file.
+            plugins: The plugins to add to the discovery file.
+        """
+        super().__init__(version=int(version))
+
+        for ptype in PluginType:
+            self[ptype] = []
+
+        for plugin_type, raw_plugins in plugins.items():
+            for raw_plugin in raw_plugins:
+                plugin_def = PluginDefinition(
+                    plugin_type,
+                    raw_plugin.pop("name"),
+                    raw_plugin.pop("namespace"),
+                    **raw_plugin,
+                )
+                self[plugin_type].append(plugin_def)
+
+    @classmethod
+    def file_version(cls, attrs):
+        """Return version of discovery file represented by attrs dictionary.
+
+        Args:
+            attrs: The attributes of the discovery file.
+
+        Returns:
+            The version of the discovery file.
+        """
+        return int(attrs.get("version", 1))

src/meltano/core/plugin_discovery_service.py

@@ -9,22 +9,25 @@
 from typing import Iterable
 
 import requests
-import yaml
+from ruamel.yaml import YAMLError
 
 import meltano
 from meltano.core import bundle
 from meltano.core.plugin.base import StandalonePlugin
 from meltano.core.project import Project
+from meltano.core.yaml import configure_yaml
 
-from .behavior.canonical import Canonical
 from .behavior.versioned import IncompatibleVersionError, Versioned
+from .discovery_file import DiscoveryFile
 from .plugin import BasePlugin, PluginDefinition, PluginRef, PluginType
 from .plugin.error import PluginNotFoundError
 from .plugin.factory import base_plugin_factory
 from .plugin.project_plugin import ProjectPlugin
 from .project_settings_service import ProjectSettingsService
 from .utils import NotFound, find_named
 
+yaml = configure_yaml()
+
 
 class DiscoveryInvalidError(Exception):
     """Occurs when the discovery.yml fails to be parsed."""
@@ -39,44 +42,6 @@ class DiscoveryUnavailableError(Exception):
 VERSION = 22
 
 
-class DiscoveryFile(Canonical):
-    """A discovery file object."""
-
-    def __init__(self, version=1, **plugins):
-        """Create a new DiscoveryFile.
-
-        Args:
-            version: The version of the discovery file.
-            plugins: The plugins to add to the discovery file.
-        """
-        super().__init__(version=int(version))
-
-        for ptype in PluginType:
-            self[ptype] = []
-
-        for plugin_type, raw_plugins in plugins.items():
-            for raw_plugin in raw_plugins:
-                plugin_def = PluginDefinition(
-                    plugin_type,
-                    raw_plugin.pop("name"),
-                    raw_plugin.pop("namespace"),
-                    **raw_plugin,
-                )
-                self[plugin_type].append(plugin_def)
-
-    @classmethod
-    def file_version(cls, attrs):
-        """Return version of discovery file represented by attrs dictionary.
-
-        Args:
-            attrs: The attributes of the discovery file.
-
-        Returns:
-            The version of the discovery file.
-        """
-        return int(attrs.get("version", 1))
-
-
 class PluginRepository(metaclass=ABCMeta):
     """A generic plugin definition repository."""
 
@@ -346,7 +311,7 @@ def load_discovery(self, discovery_file, cache=False) -> DiscoveryFile:
             DiscoveryInvalidError: If the discovery file is invalid.
         """
         try:
-            discovery_yaml = yaml.safe_load(discovery_file)
+            discovery_yaml = yaml.load(discovery_file)
 
             self._discovery_version = DiscoveryFile.file_version(discovery_yaml)
             self.ensure_compatible()
@@ -357,7 +322,7 @@ def load_discovery(self, discovery_file, cache=False) -> DiscoveryFile:
                 self.cache_discovery()
 
             return self._discovery
-        except (yaml.YAMLError, Exception) as err:
+        except (YAMLError, Exception) as err:
             raise DiscoveryInvalidError(str(err))
 
     def cache_discovery(self):
@@ -366,8 +331,6 @@ def cache_discovery(self):
             yaml.dump(
                 self._discovery,
                 cached_discovery,
-                default_flow_style=False,
-                sort_keys=False,
             )
 
     @property

src/meltano/core/transform_add_service.py

@@ -1,18 +1,29 @@
+"""Helper class for dbt package installation."""
+
 from __future__ import annotations
 
 import os
 from pathlib import Path
 
-import yaml
+from meltano.core.yaml import configure_yaml
 
 from .plugin.project_plugin import ProjectPlugin
 from .plugin.settings_service import PluginSettingsService
 from .project import Project
 from .project_plugins_service import ProjectPluginsService
 
+yaml = configure_yaml()
+
 
 class TransformAddService:
-    def __init__(self, project: Project):
+    """Helper class for adding a dbt package to the project."""
+
+    def __init__(self, project: Project) -> None:
+        """Create a new TransformAddService.
+
+        Args:
+            project: The project to add the dbt package to.
+        """
         self.project = project
 
         self.plugins_service = ProjectPluginsService(project)
@@ -28,11 +39,19 @@ def __init__(self, project: Project):
         self.packages_file = dbt_project_path.joinpath("packages.yml")
         self.dbt_project_file = dbt_project_path.joinpath("dbt_project.yml")
 
-    def add_to_packages(self, plugin: ProjectPlugin):
+    def add_to_packages(self, plugin: ProjectPlugin) -> None:
+        """Add the plugin's package to the project's `packages.yml` file.
+
+        Args:
+            plugin: The plugin to add to the project.
+
+        Raises:
+            ValueError: If the plugin is missing the git repo URL.
+        """
         if not os.path.exists(self.packages_file):
             self.packages_file.touch()
 
-        package_yaml = yaml.safe_load(self.packages_file.open()) or {"packages": []}
+        package_yaml = yaml.load(self.packages_file.open()) or {"packages": []}
 
         git_repo = plugin.pip_url
         if not git_repo:
@@ -55,12 +74,15 @@ def add_to_packages(self, plugin: ProjectPlugin):
         package_yaml["packages"].append(package_ref)
 
         with open(self.packages_file, "w") as f:
-            f.write(yaml.dump(package_yaml, default_flow_style=False, sort_keys=False))
+            yaml.dump(package_yaml, f)
 
-    def update_dbt_project(self, plugin: ProjectPlugin):
+    def update_dbt_project(self, plugin: ProjectPlugin) -> None:
         """Set transform package variables in `dbt_project.yml`.
 
         If not already present, the package name will also be added under dbt 'models'.
+
+        Args:
+            plugin: The plugin to add to the project.
         """
         settings_service = PluginSettingsService(
             self.project, plugin, plugins_service=self.plugins_service
@@ -69,7 +91,7 @@ def update_dbt_project(self, plugin: ProjectPlugin):
         package_name = settings_service.get("_package_name")
         package_vars = settings_service.get("_vars")
 
-        dbt_project_yaml = yaml.safe_load(self.dbt_project_file.open())
+        dbt_project_yaml = yaml.load(self.dbt_project_file.open())
 
         model_def = {}
 
@@ -87,6 +109,4 @@ def update_dbt_project(self, plugin: ProjectPlugin):
         dbt_project_yaml["models"][package_name] = model_def
 
         with open(self.dbt_project_file, "w") as f:
-            f.write(
-                yaml.dump(dbt_project_yaml, default_flow_style=False, sort_keys=False)
-            )
+            yaml.dump(dbt_project_yaml, f)

src/meltano/core/yaml.py

@@ -5,6 +5,7 @@
 from ruamel.yaml import YAML
 
 from meltano.core.behavior.canonical import Canonical
+from meltano.core.discovery_file import DiscoveryFile
 from meltano.core.plugin import PluginType
 from meltano.core.setting_definition import SettingKind
 
@@ -21,5 +22,6 @@ def configure_yaml() -> YAML:
     yaml.register_class(Canonical)
     yaml.register_class(PluginType)
     yaml.register_class(SettingKind)
+    yaml.register_class(DiscoveryFile)
 
     return yaml

tests/meltano/core/test_plugin_discovery_service.py

@@ -12,11 +12,7 @@
 from meltano.core import bundle
 from meltano.core.plugin import PluginType, Variant, VariantNotFoundError
 from meltano.core.plugin.project_plugin import ProjectPlugin
-from meltano.core.plugin_discovery_service import (
-    VERSION,
-    DiscoveryFile,
-    PluginNotFoundError,
-)
+from meltano.core.plugin_discovery_service import VERSION, PluginNotFoundError
 from meltano.core.project_plugins_service import PluginAlreadyAddedException
 from meltano.core.yaml import configure_yaml
 
@@ -69,7 +65,6 @@ def test_discovery_url_mock(self, subject):
     @pytest.fixture
     def discovery_yaml(self, subject):
         """Disable the discovery mock."""
-        yaml.register_class(DiscoveryFile)
         with subject.project.root_dir("discovery.yml").open("w") as discovery_yaml:
             yaml.dump(subject._discovery, discovery_yaml)