Use C-accelerated safe_load for YAML

paasta_tools/api/api.py

@@ -24,7 +24,6 @@
 import manhole
 import requests_cache
 import service_configuration_lib
-import yaml
 from pyramid.config import Configurator
 from wsgicors import CORS
 
@@ -33,6 +32,7 @@
 from paasta_tools.api import settings
 from paasta_tools.api.tweens import profiling
 from paasta_tools.api.tweens import request_logger
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.utils import load_system_paasta_config
 
 try:
@@ -233,7 +233,7 @@ def setup_clog(config_file="/nail/srv/configs/clog.yaml"):
     if clog:
         if os.path.exists(config_file):
             with open(config_file) as fp:
-                clog_config = yaml.safe_load(fp)
+                clog_config = safe_load_yaml(fp)
         else:
             # these are barebones basic configs from /nail/srv/configs/clog.yaml
             clog_config = {

paasta_tools/cli/cmds/validate.py

@@ -33,7 +33,6 @@
 from typing import Union
 
 import pytz
-import yaml
 from croniter import croniter
 from jsonschema import Draft4Validator
 from jsonschema import exceptions
@@ -63,6 +62,7 @@
 from paasta_tools.long_running_service_tools import METRICS_PROVIDER_PISCINA
 from paasta_tools.long_running_service_tools import METRICS_PROVIDER_PROMQL
 from paasta_tools.long_running_service_tools import METRICS_PROVIDER_UWSGI
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.secret_tools import get_secret_name_from_ref
 from paasta_tools.secret_tools import is_secret_ref
 from paasta_tools.secret_tools import is_shared_secret
@@ -341,7 +341,7 @@ def get_config_file_dict(file_path: str, use_ruamel: bool = False) -> Dict[Any,
                 )
                 return ruamel_loader.load(config_file)
             else:
-                return yaml.safe_load(config_file)
+                return safe_load_yaml(config_file)
         elif extension == ".json":
             return json.loads(config_file)
         else:

paasta_tools/cli/fsm/autosuggest.py

@@ -14,7 +14,7 @@
 import os
 import random
 
-import yaml
+from paasta_tools.paasta_yaml import safe_load_yaml
 
 
 def _get_smartstack_proxy_ports_from_file(root, file):
@@ -24,7 +24,7 @@ def _get_smartstack_proxy_ports_from_file(root, file):
     """
     ports = set()
     with open(os.path.join(root, file)) as f:
-        data = yaml.safe_load(f)
+        data = safe_load_yaml(f)
 
     if file.endswith("service.yaml") and "smartstack" in data:
         # Specifying this in service.yaml is old and deprecated and doesn't

paasta_tools/contrib/check_orphans.py

@@ -19,10 +19,11 @@
 from typing import Set
 from typing import Tuple
 
-import yaml
 from kazoo.client import KazooClient
 from kazoo.exceptions import NoNodeError
 
+from paasta_tools.paasta_yaml import safe_load_yaml
+
 logger = logging.getLogger("check_orphans")
 
 PREFIX = "/smartstack/global/"
@@ -40,7 +41,7 @@ class ExitCode(Enum):
 
 def get_zk_hosts(path: str) -> List[str]:
     with open(path) as f:
-        x = yaml.safe_load(f)
+        x = safe_load_yaml(f)
     return [f"{host}:{port}" for host, port in x]
 
 

paasta_tools/contrib/render_template.py

@@ -3,7 +3,7 @@
 import os
 import re
 
-import yaml
+from paasta_tools.paasta_yaml import safe_load_yaml
 
 
 def replace(s, values):
@@ -103,7 +103,7 @@ def render_values(src: str, dst: str, values: str, overwrite=True) -> None:
     config_dict = {}
     if values is not None:
         with open(values) as f:
-            config_dict = yaml.safe_load(f)
+            config_dict = safe_load_yaml(f)
     # Replace environment variables in values.yaml with environment variables
     for k, v in config_dict.items():
         config_dict[k] = re.sub(

paasta_tools/contrib/timeouts_metrics_prom.py

@@ -16,12 +16,12 @@
 """
 import os
 
-import yaml
 from prometheus_client import CollectorRegistry
 from prometheus_client import Gauge
 from prometheus_client import write_to_textfile
 from prometheus_client.metrics import MetricWrapperBase
 
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.utils import DEFAULT_SOA_DIR
 
 PROM_OUTPUT_FILE = f"{DEFAULT_SOA_DIR}/.autotune_timeouts.prom"
@@ -32,7 +32,7 @@ def read_and_write_timeouts_metrics(
     root: str, service: str, prom_metric: MetricWrapperBase
 ) -> None:
     with open(os.path.join(root, "smartstack.yaml")) as smartstack_file:
-        smartstack_yaml = yaml.safe_load(smartstack_file)
+        smartstack_yaml = safe_load_yaml(smartstack_file)
     for instance_name, info in smartstack_yaml.items():
         upstream = f"{service}.{instance_name}.egress_cluster"
         if "endpoint_timeouts" in info:

paasta_tools/envoy_tools.py

@@ -30,10 +30,10 @@
 from typing import Tuple
 
 import requests
-import yaml
 from kubernetes.client import V1Pod
 from mypy_extensions import TypedDict
 
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.utils import get_user_agent
 
 
@@ -154,7 +154,7 @@ def get_backends_from_eds(namespace: str, envoy_eds_path: str) -> List[Tuple[str
 
     if os.access(eds_file_for_namespace, os.R_OK):
         with open(eds_file_for_namespace) as f:
-            eds_yaml = yaml.safe_load(f)
+            eds_yaml = safe_load_yaml(f)
             for resource in eds_yaml.get("resources", []):
                 endpoints = resource.get("endpoints")
                 # endpoints could be None if there are no backends listed

paasta_tools/generate_authenticating_services.py

@@ -24,8 +24,7 @@
 from typing import List
 from typing import Set
 
-import yaml
-
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.utils import DEFAULT_SOA_DIR
 from paasta_tools.utils import load_system_paasta_config
 from paasta_tools.utils import write_json_configuration_file
@@ -39,7 +38,7 @@ def list_services_in_authz_config(config_path: str) -> Set[str]:
     auth_config: dict = {}
     try:
         with open(config_path) as f:
-            auth_config = yaml.safe_load(f)
+            auth_config = safe_load_yaml(f)
     except Exception as e:
         logging.warning(f"Issue loading {config_path}: {e}")
     return {

paasta_tools/paasta_yaml.py

@@ -0,0 +1,16 @@
+import yaml
+
+try:
+    from yaml import CSafeLoader as Loader
+except ImportError:  # pragma: no cover
+    from yaml import SafeLoader as Loader  # type: ignore
+
+
+def safe_load_yaml(stream):
+    """
+    yaml.safe_load() equivalent that will use a CSafeLoader if available.
+
+    This provides a significant speedup, but we allow falling back to the pure Python
+    codepath as we'd rather have things be slower than crash outright.
+    """
+    return yaml.load(stream, Loader=Loader)

paasta_tools/setup_tron_namespace.py

@@ -26,12 +26,11 @@
 from typing import Dict
 from typing import List
 
-import ruamel.yaml as yaml
-
 from paasta_tools import spark_tools
 from paasta_tools import tron_tools
 from paasta_tools.kubernetes_tools import ensure_service_account
 from paasta_tools.kubernetes_tools import KubeClient
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.tron_tools import KUBERNETES_NAMESPACE
 from paasta_tools.tron_tools import load_tron_service_config
 from paasta_tools.tron_tools import MASTER_NAMESPACE
@@ -168,7 +167,7 @@ def main():
             log.exception(f"Error while updating {MASTER_NAMESPACE}:")
 
     k8s_enabled_for_cluster = (
-        yaml.safe_load(master_config).get("k8s_options", {}).get("enabled", False)
+        safe_load_yaml(master_config).get("k8s_options", {}).get("enabled", False)
     )
     new_configs: Dict[str, str] = {}  # service -> new_config
     for service in sorted(services):

paasta_tools/tron/client.py

@@ -15,8 +15,8 @@
 from urllib.parse import urljoin
 
 import requests
-import yaml
 
+from paasta_tools.paasta_yaml import safe_load_yaml
 from paasta_tools.utils import get_user_agent
 
 
@@ -83,7 +83,7 @@ def update_namespace(self, namespace, new_config, skip_if_unchanged=True):
         current_config = self._get("/api/config", {"name": namespace, "no_header": 1})
 
         if skip_if_unchanged:
-            if yaml.safe_load(new_config) == yaml.safe_load(current_config["config"]):
+            if safe_load_yaml(new_config) == safe_load_yaml(current_config["config"]):
                 log.debug("No change in config, skipping update.")
                 return
 
@@ -112,7 +112,7 @@ def update_namespaces(
         for namespace, new_config in new_configs.items():
             current_config = current_configs.get(namespace, {})
             if skip_if_unchanged:
-                if yaml.safe_load(new_config) == yaml.safe_load(
+                if safe_load_yaml(new_config) == safe_load_yaml(
                     current_config["config"]
                 ):
                     log.debug("No change in config, skipping update.")

paasta_tools/tron_tools.py

@@ -33,6 +33,7 @@
 from service_configuration_lib.spark_config import SparkConfBuilder
 
 from paasta_tools.mesos_tools import mesos_services_running_here
+from paasta_tools.paasta_yaml import safe_load_yaml
 
 try:
     from yaml.cyaml import CSafeDumper as Dumper
@@ -1338,7 +1339,7 @@ def validate_complete_config(
     # given that the kill-switch will affect PaaSTA's setup_tron_namespace script (we're
     # not reading the kill-switch in Tron since it's not easily accessible at the point
     # at which we'd like to fallback to Mesos if toggled)
-    master_config = yaml.safe_load(
+    master_config = safe_load_yaml(
         create_complete_master_config(cluster=cluster, soa_dir=soa_dir)
     )
     k8s_enabled_for_cluster = master_config.get("k8s_options", {}).get("enabled", False)

paasta_tools/utils.py

@@ -83,6 +83,7 @@
 from service_configuration_lib import read_service_configuration
 
 import paasta_tools.cli.fsm
+from paasta_tools.paasta_yaml import safe_load_yaml
 
 
 # DO NOT CHANGE SPACER, UNLESS YOU'RE PREPARED TO CHANGE ALL INSTANCES
@@ -4252,7 +4253,7 @@ def write_yaml_configuration_file(
     :param bool check_existing: if existing file already matches config, do not overwrite
     """
     if check_existing:
-        previous_config = maybe_load_previous_config(filename, yaml.safe_load)
+        previous_config = maybe_load_previous_config(filename, safe_load_yaml)
         if previous_config and previous_config == configuration:
             return
 

tests/cli/fsm/test_autosuggest.py

@@ -24,9 +24,9 @@ class TestGetSmartstackProxyPortFromFile:
     def test_multiple_stanzas_per_file(self, mock_read_etc_services):
         with mock.patch("builtins.open", autospec=True):
             with mock.patch(
-                "paasta_tools.cli.fsm.autosuggest.yaml", autospec=True
+                "paasta_tools.cli.fsm.autosuggest.safe_load_yaml", autospec=True
             ) as mock_yaml:
-                mock_yaml.safe_load.return_value = {
+                mock_yaml.return_value = {
                     "main": {"proxy_port": 1},
                     "foo": {"proxy_port": 2},
                 }

tests/test_envoy_tools.py

@@ -319,7 +319,7 @@ def test_are_services_up_on_port_missing_backend(
 
     @mock.patch("paasta_tools.envoy_tools.open", autospec=False)
     @mock.patch("paasta_tools.envoy_tools.os.access", autospec=True)
-    @mock.patch("paasta_tools.envoy_tools.yaml.safe_load", autospec=True)
+    @mock.patch("paasta_tools.envoy_tools.safe_load_yaml", autospec=True)
     def test_get_backends_from_eds(self, mock_yaml, mock_os_access, mock_open):
 
         mock_yaml.return_value = {"resources": [{"endpoints": None}]}