Skip to content

Commit e6bf483

Browse files
noelmcloughlinnoelmcloughlin
authored
fix(validation_context): optimize logic around json schema
1 parent 08d487a commit e6bf483

2 files changed

Lines changed: 844 additions & 8 deletions

File tree

‎packages/linkml/src/linkml/validator/validation_context.py‎

Lines changed: 90 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,55 @@
11
import json
22
import os
3+
from collections import OrderedDict
34
from functools import lru_cache
45

56
import jsonschema
67
from jsonschema.protocols import Validator
78

89
from linkml.generators import JsonSchemaGenerator, PydanticGenerator
10+
from linkml.generators.jsonschemagen import JsonSchema
911
from linkml.utils.datautils import infer_root_class
1012
from linkml_runtime import SchemaView
1113
from linkml_runtime.linkml_model import SchemaDefinition
14+
from linkml_runtime.utils.formatutils import camelcase
15+
16+
# Module-level bounded cache of generated JSON Schemas.
17+
# See https://github.com/linkml/linkml/pull/3430 for discussion.
18+
#
19+
# Key: (id(schema), include_range_class_descendants).
20+
# Using id(schema) as the discriminator avoids two failure modes of any
21+
# metadata-based key (id/name/version): (a) workflows that load multiple
22+
# distinct schemas under a placeholder id like "http://example.org/default"
23+
# (test_compliance.py); (b) workflows that merge overlays into a base schema
24+
# via `--include`, producing a logically different schema that inherits the
25+
# base's id/name/version. Distinct Python objects always get distinct keys.
26+
#
27+
# To guard against CPython recycling id() after a schema is GC'd (aliasing
28+
# the new object onto a stale cache entry), the `_schema_pins` dict has
29+
# a strong reference to each cached schema. The ref drops only when the
30+
# entry is evicted by the bounded LRU policy - allowing id reuse, but
31+
# the cache no longer has any stale entry to alias against.
32+
# See: https://en.wikipedia.org/wiki/Cache_replacement_policies
33+
#
34+
# The root's additionalProperties is overridden at retrieval time, so cached
35+
# $defs are reusable regardless of closed.
36+
_JSON_SCHEMA_CACHE_MAXSIZE = 32
37+
_json_schema_cache: OrderedDict[tuple, JsonSchema] = OrderedDict()
38+
_schema_pins: dict[tuple, SchemaDefinition] = {}
39+
40+
41+
def _make_cache_key(
42+
schema: SchemaDefinition,
43+
include_range_class_descendants: bool,
44+
) -> tuple:
45+
return (id(schema), include_range_class_descendants)
46+
47+
48+
# Keys that JsonSchemaGenerator's start_schema sets on the root and that the
49+
# per-class merge step in handle_class does NOT overwrite. These are
50+
# schema-level: they describe the schema as a whole, not whichever class
51+
# happens to have warmed the cache. Safe to inherit from the cached root.
52+
_ROOT_METADATA_KEYS: tuple = ("$schema", "$id", "metamodel_version", "version", "title", "type")
1253

1354

1455
class ValidationContext:
@@ -42,14 +83,55 @@ def json_schema_validator(
4283
json_schema = json.load(json_schema_file)
4384
else:
4485
not_closed = not closed
45-
jsonschema_gen = JsonSchemaGenerator(
46-
schema=self._schema,
47-
mergeimports=True,
48-
top_class=self._target_class,
49-
not_closed=not_closed,
50-
include_range_class_descendants=include_range_class_descendants,
51-
)
52-
json_schema = jsonschema_gen.generate()
86+
cache_key = _make_cache_key(self._schema, include_range_class_descendants)
87+
88+
if cache_key not in _json_schema_cache:
89+
# First call: generate and cache entire schema (full generation cost)
90+
jsonschema_gen = JsonSchemaGenerator(
91+
schema=self._schema,
92+
mergeimports=True,
93+
top_class=self._target_class,
94+
not_closed=not_closed,
95+
include_range_class_descendants=include_range_class_descendants,
96+
)
97+
json_schema = jsonschema_gen.generate()
98+
_json_schema_cache[cache_key] = json_schema
99+
# Pin the schema so id() cannot be recycled while this entry exists.
100+
_schema_pins[cache_key] = self._schema
101+
while len(_json_schema_cache) > _JSON_SCHEMA_CACHE_MAXSIZE:
102+
evicted_key, _ = _json_schema_cache.popitem(last=False)
103+
_schema_pins.pop(evicted_key, None)
104+
else:
105+
# Subsequent calls: reuse cached $defs, rebuild the root by re-doing
106+
# JsonSchemaGenerator's "merge top class into root" step for a new target.
107+
# Wrapping with $ref would inherit the hardcoded additionalProperties=False
108+
# from $defs[X] (see jsonschemagen.py handle_class) and diverge from a
109+
# freshly-generated validator when closed=False.
110+
_json_schema_cache.move_to_end(cache_key) # mark as MRU (most recently used)
111+
cached = _json_schema_cache[cache_key]
112+
# $defs keys are camelCased by JsonSchemaGenerator when preserve_names=False
113+
# (the mode used here), so look up under the canonical name.
114+
defs_key = camelcase(self._target_class)
115+
defs_class = cached["$defs"].get(defs_key, {})
116+
# Build the root in three layers:
117+
# 1) Schema-level metadata inherited from cached (immune to which class
118+
# happened to warm the cache).
119+
# 2) Class-specific keys from $defs[target_class] (properties, required,
120+
# description, if/then/else, allOf, plus any extension-hook additions).
121+
# Exclude $schema/$id/$defs/additionalProperties (handled separately)
122+
# and the metadata keys (already inherited from cached).
123+
# 3) additionalProperties explicitly set to not_closed.
124+
root = {k: cached[k] for k in _ROOT_METADATA_KEYS if k in cached}
125+
root["$defs"] = cached["$defs"]
126+
root.update(
127+
{
128+
k: v
129+
for k, v in defs_class.items()
130+
if k not in _ROOT_METADATA_KEYS and k not in ("$schema", "$id", "$defs", "additionalProperties")
131+
}
132+
)
133+
root["additionalProperties"] = not_closed
134+
json_schema = JsonSchema(root)
53135

54136
validator_cls = jsonschema.validators.validator_for(json_schema, default=jsonschema.Draft7Validator)
55137
return validator_cls(json_schema, format_checker=validator_cls.FORMAT_CHECKER)

0 commit comments

Comments
 (0)