wehub-resource-sync
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
from typing import NamedTuple
|
|
|
|
from graphql.error import GraphQLError
|
|
from graphql.execution import ExecutionResult
|
|
from graphql.language.ast import DocumentNode, FieldNode
|
|
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_SERVER_GRAPHQL_MAX_ALIASES,
|
|
MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS,
|
|
)
|
|
|
|
_MAX_DEPTH = 10
|
|
_MAX_SELECTIONS = 1000
|
|
|
|
|
|
class QueryInfo(NamedTuple):
|
|
root_fields: int
|
|
max_aliases: int
|
|
|
|
|
|
def scan_query(ast_node: DocumentNode) -> QueryInfo:
|
|
"""
|
|
Scan a GraphQL query and return its information.
|
|
"""
|
|
root_fields = 0
|
|
max_aliases = 0
|
|
total_selections = 0
|
|
|
|
for definition in ast_node.definitions:
|
|
if selection_set := getattr(definition, "selection_set", None):
|
|
stack = [(selection_set, 1)]
|
|
while stack:
|
|
selection_set, depth = stack.pop()
|
|
|
|
# check current level depth
|
|
if depth > _MAX_DEPTH:
|
|
raise GraphQLError(f"Query exceeds maximum depth of {_MAX_DEPTH}")
|
|
|
|
selections = getattr(selection_set, "selections", [])
|
|
|
|
# check current level aliases
|
|
current_aliases = 0
|
|
for selection in selections:
|
|
if isinstance(selection, FieldNode):
|
|
if depth == 1:
|
|
root_fields += 1
|
|
if selection.alias:
|
|
current_aliases += 1
|
|
if selection.selection_set:
|
|
stack.append((selection.selection_set, depth + 1))
|
|
total_selections += 1
|
|
if total_selections > _MAX_SELECTIONS:
|
|
raise GraphQLError(
|
|
f"Query exceeds maximum total selections of {_MAX_SELECTIONS}"
|
|
)
|
|
max_aliases = max(max_aliases, current_aliases)
|
|
|
|
return QueryInfo(root_fields, max_aliases)
|
|
|
|
|
|
def check_query_safety(ast_node: DocumentNode) -> ExecutionResult | None:
|
|
try:
|
|
query_info = scan_query(ast_node)
|
|
except GraphQLError as e:
|
|
return ExecutionResult(
|
|
data=None,
|
|
errors=[e],
|
|
)
|
|
|
|
if query_info.root_fields > MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS.get():
|
|
msg = "root fields"
|
|
env_var = MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS
|
|
value = query_info.root_fields
|
|
elif query_info.max_aliases > MLFLOW_SERVER_GRAPHQL_MAX_ALIASES.get():
|
|
msg = "aliases"
|
|
env_var = MLFLOW_SERVER_GRAPHQL_MAX_ALIASES
|
|
value = query_info.max_aliases
|
|
else:
|
|
return None
|
|
return ExecutionResult(
|
|
data=None,
|
|
errors=[
|
|
GraphQLError(
|
|
f"GraphQL queries should have at most {env_var.get()} {msg}, "
|
|
f"got {value} {msg}. To increase the limit, set the "
|
|
f"{env_var.name} environment variable."
|
|
)
|
|
],
|
|
)
|