wehub-resource-sync
333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""
|
|
JSONPath utilities for navigating and manipulating nested JSON structures.
|
|
|
|
This module provides a simplified JSONPath-like implementation without adding
|
|
external dependencies to MLflow. Instead of using a full JSONPath library,
|
|
we implement a lightweight subset focused on trace data navigation using
|
|
dot notation with wildcard support.
|
|
|
|
The implementation supports:
|
|
- Dot notation path traversal (e.g., "info.trace_id")
|
|
- Wildcard expansion (e.g., "info.assessments.*")
|
|
- Array/list navigation with numeric indices
|
|
- Structure-preserving filtering
|
|
- Path validation with helpful error messages
|
|
|
|
This approach keeps MLflow dependencies minimal while providing the essential
|
|
functionality needed for trace field selection and data manipulation.
|
|
|
|
Note: This is NOT a complete JSONPath implementation. It's a custom solution
|
|
tailored specifically for MLflow trace data structures.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
|
|
def split_path_respecting_backticks(path: str) -> list[str]:
|
|
"""
|
|
Split path on dots, but keep backticked segments intact.
|
|
|
|
Args:
|
|
path: Path string like 'info.tags.`mlflow.traceName`'
|
|
|
|
Returns:
|
|
List of path segments, e.g., ['info', 'tags', 'mlflow.traceName']
|
|
"""
|
|
parts = []
|
|
i = 0
|
|
current = ""
|
|
|
|
while i < len(path):
|
|
if i < len(path) and path[i] == "`":
|
|
# Start of backticked segment - read until closing backtick
|
|
i += 1 # Skip opening backtick
|
|
while i < len(path) and path[i] != "`":
|
|
current += path[i]
|
|
i += 1
|
|
if i < len(path):
|
|
i += 1 # Skip closing backtick
|
|
elif path[i] == ".":
|
|
if current:
|
|
parts.append(current)
|
|
current = ""
|
|
i += 1
|
|
else:
|
|
current += path[i]
|
|
i += 1
|
|
|
|
if current:
|
|
parts.append(current)
|
|
|
|
return parts
|
|
|
|
|
|
def jsonpath_extract_values(obj: dict[str, Any], path: str) -> list[Any]:
|
|
"""
|
|
Extract values from nested dict using JSONPath-like dot notation with * wildcard support.
|
|
|
|
Supports backtick escaping for field names containing dots:
|
|
'info.tags.`mlflow.traceName`' - treats 'mlflow.traceName' as a single field
|
|
|
|
Args:
|
|
obj: The dictionary/object to traverse
|
|
path: Dot-separated path like 'info.trace_id' or 'data.spans.*.name'
|
|
Can use backticks for fields with dots: 'info.tags.`mlflow.traceName`'
|
|
|
|
Returns:
|
|
List of values found at the path. Returns empty list if path not found.
|
|
|
|
Examples:
|
|
>>> data = {"info": {"trace_id": "tr-123", "status": "OK"}}
|
|
>>> jsonpath_extract_values(data, "info.trace_id")
|
|
['tr-123']
|
|
>>> jsonpath_extract_values(data, "info.*")
|
|
['tr-123', 'OK']
|
|
>>> data = {"tags": {"mlflow.traceName": "test"}}
|
|
>>> jsonpath_extract_values(data, "tags.`mlflow.traceName`")
|
|
['test']
|
|
"""
|
|
parts = split_path_respecting_backticks(path)
|
|
|
|
def traverse(current, parts_remaining):
|
|
if not parts_remaining:
|
|
return [current]
|
|
|
|
part = parts_remaining[0]
|
|
rest = parts_remaining[1:]
|
|
|
|
if part == "*":
|
|
# Wildcard - expand all keys at this level
|
|
if isinstance(current, dict):
|
|
results = []
|
|
for key, value in current.items():
|
|
results.extend(traverse(value, rest))
|
|
return results
|
|
elif isinstance(current, list):
|
|
results = []
|
|
for item in current:
|
|
results.extend(traverse(item, rest))
|
|
return results
|
|
else:
|
|
return []
|
|
else:
|
|
# Regular key
|
|
if isinstance(current, dict) and part in current:
|
|
return traverse(current[part], rest)
|
|
else:
|
|
return []
|
|
|
|
return traverse(obj, parts)
|
|
|
|
|
|
def filter_json_by_fields(data: dict[str, Any], field_paths: list[str]) -> dict[str, Any]:
|
|
"""
|
|
Filter a JSON dict to only include fields specified by the field paths.
|
|
Expands wildcards but preserves original JSON structure.
|
|
|
|
Args:
|
|
data: Original JSON dictionary
|
|
field_paths: List of dot-notation paths like ['info.trace_id', 'info.assessments.*']
|
|
|
|
Returns:
|
|
Filtered dictionary with original structure preserved
|
|
"""
|
|
result = {}
|
|
|
|
# Collect all actual paths by expanding wildcards
|
|
expanded_paths = set()
|
|
for field_path in field_paths:
|
|
if "*" in field_path:
|
|
# Find all actual paths that match this wildcard pattern
|
|
matching_paths = find_matching_paths(data, field_path)
|
|
expanded_paths.update(matching_paths)
|
|
else:
|
|
# Direct path
|
|
expanded_paths.add(field_path)
|
|
|
|
# Build the result by including only the specified paths
|
|
for path in expanded_paths:
|
|
parts = split_path_respecting_backticks(path)
|
|
set_nested_value(result, parts, get_nested_value_safe(data, parts))
|
|
|
|
return result
|
|
|
|
|
|
def find_matching_paths(data: dict[str, Any], wildcard_path: str) -> list[str]:
|
|
"""Find all actual paths in data that match a wildcard pattern."""
|
|
parts = split_path_respecting_backticks(wildcard_path)
|
|
|
|
def find_paths(current_data, current_parts, current_path=""):
|
|
if not current_parts:
|
|
return [current_path.lstrip(".")]
|
|
|
|
part = current_parts[0]
|
|
remaining = current_parts[1:]
|
|
|
|
if part == "*":
|
|
paths = []
|
|
if isinstance(current_data, dict):
|
|
for key in current_data.keys():
|
|
new_path = f"{current_path}.{key}"
|
|
paths.extend(find_paths(current_data[key], remaining, new_path))
|
|
elif isinstance(current_data, list):
|
|
for i, item in enumerate(current_data):
|
|
new_path = f"{current_path}.{i}"
|
|
paths.extend(find_paths(item, remaining, new_path))
|
|
return paths
|
|
else:
|
|
if isinstance(current_data, dict) and part in current_data:
|
|
new_path = f"{current_path}.{part}"
|
|
return find_paths(current_data[part], remaining, new_path)
|
|
return []
|
|
|
|
return find_paths(data, parts)
|
|
|
|
|
|
def get_nested_value_safe(data: dict[str, Any], parts: list[str]) -> Any | None:
|
|
"""Safely get nested value, returning None if path doesn't exist."""
|
|
current = data
|
|
for part in parts:
|
|
if isinstance(current, dict) and part in current:
|
|
current = current[part]
|
|
elif isinstance(current, list) and part.isdigit() and int(part) < len(current):
|
|
current = current[int(part)]
|
|
else:
|
|
return None
|
|
return current
|
|
|
|
|
|
def set_nested_value(data: dict[str, Any], parts: list[str], value: Any) -> None:
|
|
"""Set a nested value in a dictionary, creating intermediate dicts/lists as needed."""
|
|
if value is None:
|
|
return
|
|
|
|
current = data
|
|
for i, part in enumerate(parts[:-1]):
|
|
if part.isdigit() and isinstance(current, list):
|
|
# Handle array index
|
|
idx = int(part)
|
|
while len(current) <= idx:
|
|
current.append({})
|
|
current = current[idx]
|
|
else:
|
|
# Handle object key
|
|
if not isinstance(current, dict):
|
|
return # Can't set object key on non-dict
|
|
if part not in current:
|
|
# Look ahead to see if next part is a number (array index)
|
|
next_part = parts[i + 1] if i + 1 < len(parts) else None
|
|
if next_part and next_part.isdigit():
|
|
current[part] = []
|
|
else:
|
|
current[part] = {}
|
|
current = current[part]
|
|
|
|
if parts:
|
|
final_part = parts[-1]
|
|
if final_part.isdigit() and isinstance(current, list):
|
|
# Extend list if needed
|
|
idx = int(final_part)
|
|
while len(current) <= idx:
|
|
current.append(None)
|
|
current[idx] = value
|
|
elif isinstance(current, dict):
|
|
current[final_part] = value
|
|
|
|
|
|
def validate_field_paths(
|
|
field_paths: list[str], sample_data: dict[str, Any], verbose: bool = False
|
|
) -> None:
|
|
"""Validate that field paths exist in the data structure.
|
|
|
|
Args:
|
|
field_paths: List of field paths to validate
|
|
sample_data: Sample data to validate against
|
|
verbose: If True, show all available fields instead of truncated list
|
|
"""
|
|
invalid_paths = []
|
|
|
|
for path in field_paths:
|
|
# Skip validation for paths with wildcards - they'll be expanded later
|
|
if "*" in path:
|
|
continue
|
|
|
|
# Test if the path exists by trying to extract values
|
|
values = jsonpath_extract_values(sample_data, path)
|
|
if not values: # Empty list means path doesn't exist
|
|
invalid_paths.append(path)
|
|
|
|
if invalid_paths:
|
|
available_fields = get_available_field_suggestions(sample_data)
|
|
|
|
# Create a nice error message
|
|
error_msg = "❌ Invalid field path(s):\n"
|
|
for path in invalid_paths:
|
|
error_msg += f" • {path}\n"
|
|
|
|
error_msg += "\n💡 Use dot notation to specify nested fields:"
|
|
error_msg += "\n Examples: info.trace_id, info.state, info.assessments.*"
|
|
|
|
if available_fields:
|
|
error_msg += "\n\n📋 Available fields in this data:\n"
|
|
|
|
if verbose:
|
|
# In verbose mode, show ALL available fields organized by category
|
|
info_fields = [f for f in available_fields if f.startswith("info.")]
|
|
data_fields = [f for f in available_fields if f.startswith("data.")]
|
|
|
|
if info_fields:
|
|
error_msg += " Info fields:\n"
|
|
for field in sorted(info_fields):
|
|
error_msg += f" • {field}\n"
|
|
|
|
if data_fields:
|
|
error_msg += " Data fields:\n"
|
|
for field in sorted(data_fields):
|
|
error_msg += f" • {field}\n"
|
|
else:
|
|
# Non-verbose mode: show truncated list
|
|
# Group by top-level key for better readability
|
|
info_fields = [f for f in available_fields if f.startswith("info.")]
|
|
data_fields = [f for f in available_fields if f.startswith("data.")]
|
|
|
|
if info_fields:
|
|
error_msg += f" info.*: {', '.join(info_fields[:8])}"
|
|
if len(info_fields) > 8:
|
|
error_msg += f", ... (+{len(info_fields) - 8} more)"
|
|
error_msg += "\n"
|
|
|
|
if data_fields:
|
|
error_msg += f" data.*: {', '.join(data_fields[:5])}"
|
|
if len(data_fields) > 5:
|
|
error_msg += f", ... (+{len(data_fields) - 5} more)"
|
|
error_msg += "\n"
|
|
|
|
error_msg += "\n💡 Tip: Use --verbose flag to see all available fields"
|
|
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
def get_available_field_suggestions(data: dict[str, Any], prefix: str = "") -> list[str]:
|
|
"""Get a list of available field paths for suggestions."""
|
|
paths = []
|
|
|
|
def collect_paths(obj, current_path=""):
|
|
if isinstance(obj, dict):
|
|
for key, value in obj.items():
|
|
path = f"{current_path}.{key}" if current_path else key
|
|
paths.append(path)
|
|
# Only go 2 levels deep for suggestions to keep it manageable
|
|
if current_path.count(".") < 2:
|
|
collect_paths(value, path)
|
|
elif isinstance(obj, list) and obj:
|
|
# Show array notation but don't expand all indices
|
|
path = f"{current_path}.*" if current_path else "*"
|
|
if path not in paths:
|
|
paths.append(path)
|
|
# Sample first item if it's an object
|
|
if isinstance(obj[0], dict):
|
|
collect_paths(obj[0], f"{current_path}.*" if current_path else "*")
|
|
|
|
collect_paths(data, prefix)
|
|
return sorted(set(paths))
|