From b251702af1c1d54854d98d593c444ce0a6b33f97 Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@microsoft.com>
Date: Thu, 6 Jul 2023 20:26:04 +0000
Subject: [PATCH] Schema resolution fixups for new version of jsonschema with
 referencing lib.

---
 .../config/schemas/config_schemas.py          | 29 +++++++++++++++++--
 mlos_bench/setup.py                           |  2 +-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/mlos_bench/mlos_bench/config/schemas/config_schemas.py b/mlos_bench/mlos_bench/config/schemas/config_schemas.py
index 61055cd549e..03fcdb564e2 100644
--- a/mlos_bench/mlos_bench/config/schemas/config_schemas.py
+++ b/mlos_bench/mlos_bench/config/schemas/config_schemas.py
@@ -14,6 +14,9 @@
 import json         # schema files are pure json - no comments
 import jsonschema
 
+from referencing import Registry, Resource
+from referencing.jsonschema import DRAFT202012
+
 from mlos_bench.util import path_join
 
 _LOG = logging.getLogger(__name__)
@@ -39,6 +42,7 @@ class SchemaStore(Mapping):
 
     # A class member mapping of schema id to schema object.
     _SCHEMA_STORE: Dict[str, dict] = {}
+    _REGISTRY: Registry = Registry()
 
     def __len__(self) -> int:
         return self._SCHEMA_STORE.__len__()
@@ -55,6 +59,8 @@ def __getitem__(self, key: str) -> dict:
     @classmethod
     def _load_schemas(cls) -> None:
         """Loads all schemas and subschemas into the schema store for the validator to reference."""
+        if cls._SCHEMA_STORE:
+            return
         for root, _, files in walk(CONFIG_SCHEMA_DIR):
             for file_name in files:
                 if not file_name.endswith(".json"):
@@ -70,6 +76,23 @@ def _load_schemas(cls) -> None:
                     assert schema["$id"] not in cls._SCHEMA_STORE
                     cls._SCHEMA_STORE[schema["$id"]] = schema
 
+    @classmethod
+    def _load_registry(cls) -> None:
+        """Also store them in a Registry object for referencing by recent versions of jsonschema."""
+        if not cls._SCHEMA_STORE:
+            cls._load_schemas()
+        cls._REGISTRY = Registry().with_resources([
+            (url, Resource.from_contents(schema, default_specification=DRAFT202012))
+            for url, schema in cls._SCHEMA_STORE.items()
+        ])
+
+    @property
+    def registry(self) -> Registry:
+        """Returns a Registry object with all the schemas loaded."""
+        if not self._REGISTRY:
+            self._load_registry()
+        return self._REGISTRY
+
 
 SCHEMA_STORE = SchemaStore()
 
@@ -112,5 +135,7 @@ def validate(self, config: dict) -> None:
         if _SKIP_VALIDATION:
             _LOG.warning("%s is set - skip schema validation", _VALIDATION_ENV_FLAG)
         else:
-            resolver: jsonschema.RefResolver = jsonschema.RefResolver.from_schema(self.schema, store=SCHEMA_STORE)
-            jsonschema.validate(instance=config, schema=self.schema, resolver=resolver)
+            jsonschema.Draft202012Validator(
+                schema=self.schema,
+                registry=SCHEMA_STORE.registry,     # type: ignore[call-arg]
+            ).validate(config)
diff --git a/mlos_bench/setup.py b/mlos_bench/setup.py
index 3a0c3073dd8..9100dcff76f 100644
--- a/mlos_bench/setup.py
+++ b/mlos_bench/setup.py
@@ -76,7 +76,7 @@
         'mlos-core==' + _VERSION,
         'requests',
         'json5',
-        'jsonschema',
+        'jsonschema>=4.18.0', 'referencing>=0.29.1',
         'importlib_resources;python_version<"3.10"',
     ] + extra_requires['storage-sql-sqlite'],   # NOTE: For now sqlite is a fallback storage backend, so we always install it.
     extras_require=extra_requires,