wehub-resource-sync
226 lines
8.7 KiB
Python
226 lines
8.7 KiB
Python
import ast
|
|
|
|
from autogeneration_utils import (
|
|
DUMMY_FIELD,
|
|
INDENT,
|
|
INDENT2,
|
|
SCHEMA_EXTENSION,
|
|
SCHEMA_EXTENSION_MODULE,
|
|
get_descriptor_full_pascal_name,
|
|
get_method_name,
|
|
method_descriptor_to_generated_pb2_file_name,
|
|
)
|
|
from google.protobuf.descriptor import FieldDescriptor
|
|
from string_utils import camel_to_snake, snake_to_pascal
|
|
|
|
# Mapping from proto descriptor type to graphene object type.
|
|
PROTO_TO_GRAPHENE_TYPE = {
|
|
FieldDescriptor.TYPE_BOOL: "graphene.Boolean",
|
|
FieldDescriptor.TYPE_FLOAT: "graphene.Float",
|
|
FieldDescriptor.TYPE_INT32: "graphene.Int",
|
|
FieldDescriptor.TYPE_INT64: "LongString",
|
|
FieldDescriptor.TYPE_STRING: "graphene.String",
|
|
FieldDescriptor.TYPE_DOUBLE: "graphene.Float",
|
|
FieldDescriptor.TYPE_UINT32: "graphene.Int",
|
|
FieldDescriptor.TYPE_UINT64: "LongString",
|
|
FieldDescriptor.TYPE_SINT32: "graphene.Int",
|
|
FieldDescriptor.TYPE_SINT64: "LongString",
|
|
FieldDescriptor.TYPE_BYTES: "graphene.String",
|
|
FieldDescriptor.TYPE_FIXED32: "graphene.Int",
|
|
FieldDescriptor.TYPE_FIXED64: "LongString",
|
|
FieldDescriptor.TYPE_SFIXED32: "graphene.Int",
|
|
FieldDescriptor.TYPE_SFIXED64: "LongString",
|
|
FieldDescriptor.TYPE_ENUM: "graphene.Enum",
|
|
}
|
|
|
|
"""
|
|
Based on graphql_schema_extensions.py, constructs a map from the name of the
|
|
extended class to the name of the extending class.
|
|
For example
|
|
class AutogenExtension(OriginalAutogen)
|
|
would give us {"OriginalAutogen": "AutogenExtension"}
|
|
"""
|
|
|
|
|
|
class ClassInheritanceVisitor(ast.NodeVisitor):
|
|
def __init__(self):
|
|
self.inheritance_map = {}
|
|
|
|
def visit_ClassDef(self, node):
|
|
for base in node.bases:
|
|
if isinstance(base, ast.Name): # Direct superclass
|
|
if base.id in self.inheritance_map:
|
|
raise Exception(
|
|
f"{base.id} is being extended more than once in {SCHEMA_EXTENSION}. "
|
|
+ "A GraphQL schema class should not be extended more than once."
|
|
)
|
|
self.inheritance_map[base.id] = node.name
|
|
self.generic_visit(node)
|
|
|
|
|
|
def get_manual_extensions():
|
|
with open(SCHEMA_EXTENSION) as file:
|
|
file_content = file.read()
|
|
|
|
parsed_content = ast.parse(file_content)
|
|
visitor = ClassInheritanceVisitor()
|
|
visitor.visit(parsed_content)
|
|
|
|
return visitor.inheritance_map
|
|
|
|
|
|
# The resulting map
|
|
EXTENDED_TO_EXTENDING = get_manual_extensions()
|
|
|
|
"""
|
|
Given the GenerateSchemaState, generate the whole schema with Graphene.
|
|
"""
|
|
|
|
|
|
def generate_schema(state):
|
|
schema_builder = ""
|
|
schema_builder += "# GENERATED FILE. PLEASE DON'T MODIFY.\n"
|
|
schema_builder += "# Run uv run ./dev/proto_to_graphql/code_generator.py to regenerate.\n"
|
|
schema_builder += "import graphene\n"
|
|
schema_builder += "import mlflow\n"
|
|
schema_builder += "from mlflow.server.graphql.graphql_custom_scalars import LongString\n"
|
|
schema_builder += "from mlflow.server.graphql.graphql_errors import ApiError\n"
|
|
schema_builder += "from mlflow.utils.proto_json_utils import parse_dict\n"
|
|
schema_builder += "\n"
|
|
|
|
for enum in sorted(state.enums, key=lambda item: item.full_name):
|
|
pascal_class_name = snake_to_pascal(get_descriptor_full_pascal_name(enum))
|
|
schema_builder += f"\nclass {pascal_class_name}(graphene.Enum):"
|
|
for i in range(len(enum.values)):
|
|
value = enum.values[i]
|
|
# enum indices start from 1
|
|
schema_builder += f"""\n{INDENT}{value.name} = {i + 1}"""
|
|
schema_builder += "\n\n"
|
|
|
|
for type in state.types:
|
|
pascal_class_name = snake_to_pascal(get_descriptor_full_pascal_name(type))
|
|
schema_builder += f"\nclass {pascal_class_name}(graphene.ObjectType):"
|
|
for field in type.fields:
|
|
graphene_type = get_graphene_type_for_field(field, False)
|
|
schema_builder += f"\n{INDENT}{camel_to_snake(field.name)} = {graphene_type}"
|
|
|
|
if type in state.outputs:
|
|
schema_builder += f"\n{INDENT}apiError = graphene.Field(ApiError)"
|
|
|
|
if len(type.fields) == 0:
|
|
schema_builder += f"\n{INDENT}{DUMMY_FIELD}"
|
|
|
|
schema_builder += "\n\n"
|
|
|
|
for input in state.inputs:
|
|
pascal_class_name = snake_to_pascal(get_descriptor_full_pascal_name(input)) + "Input"
|
|
schema_builder += f"\nclass {pascal_class_name}(graphene.InputObjectType):"
|
|
for field in input.fields:
|
|
graphene_type = get_graphene_type_for_field(field, True)
|
|
schema_builder += f"\n{INDENT}{camel_to_snake(field.name)} = {graphene_type}"
|
|
if len(input.fields) == 0:
|
|
schema_builder += f"\n{INDENT}{DUMMY_FIELD}"
|
|
|
|
schema_builder += "\n\n"
|
|
|
|
schema_builder += "\nclass QueryType(graphene.ObjectType):"
|
|
|
|
if len(state.queries) == 0:
|
|
schema_builder += f"\n{INDENT}pass"
|
|
|
|
for query in sorted(state.queries, key=lambda item: item.name):
|
|
schema_builder += proto_method_to_graphql_operation(query)
|
|
|
|
schema_builder += "\n"
|
|
|
|
for query in sorted(state.queries, key=lambda item: item.name):
|
|
schema_builder += generate_resolver_function(query)
|
|
|
|
schema_builder += "\n"
|
|
schema_builder += "\nclass MutationType(graphene.ObjectType):"
|
|
|
|
if len(state.mutations) == 0:
|
|
schema_builder += f"\n{INDENT}pass"
|
|
|
|
for mutation in sorted(state.mutations, key=lambda item: item.name):
|
|
schema_builder += proto_method_to_graphql_operation(mutation)
|
|
|
|
schema_builder += "\n"
|
|
|
|
for mutation in sorted(state.mutations, key=lambda item: item.name):
|
|
schema_builder += generate_resolver_function(mutation)
|
|
|
|
return schema_builder
|
|
|
|
|
|
def apply_schema_extension(referenced_class_name):
|
|
if referenced_class_name in EXTENDED_TO_EXTENDING:
|
|
# Using dotted module path as pointed out in the linked GitHub issue.r
|
|
# This is an undocumented feature of Graphene.
|
|
# https://github.com/graphql-python/graphene/issues/110#issuecomment-1219737639
|
|
return f"'{SCHEMA_EXTENSION_MODULE}.{EXTENDED_TO_EXTENDING[referenced_class_name]}'"
|
|
else:
|
|
return referenced_class_name
|
|
|
|
|
|
def get_graphene_type_for_field(field, is_input):
|
|
if field.type == FieldDescriptor.TYPE_ENUM:
|
|
referenced_class_name = apply_schema_extension(
|
|
get_descriptor_full_pascal_name(field.enum_type)
|
|
)
|
|
if field.label == FieldDescriptor.LABEL_REPEATED:
|
|
return f"graphene.List(graphene.NonNull({referenced_class_name}))"
|
|
else:
|
|
return f"graphene.Field({referenced_class_name})"
|
|
elif field.type in (FieldDescriptor.TYPE_GROUP, FieldDescriptor.TYPE_MESSAGE):
|
|
if is_input:
|
|
referenced_class_name = apply_schema_extension(
|
|
f"{get_descriptor_full_pascal_name(field.message_type)}Input"
|
|
)
|
|
field_type_base = f"graphene.InputField({referenced_class_name})"
|
|
else:
|
|
referenced_class_name = apply_schema_extension(
|
|
get_descriptor_full_pascal_name(field.message_type)
|
|
)
|
|
field_type_base = f"graphene.Field({referenced_class_name})"
|
|
if field.label == FieldDescriptor.LABEL_REPEATED:
|
|
return f"graphene.List(graphene.NonNull({referenced_class_name}))"
|
|
else:
|
|
return field_type_base
|
|
else:
|
|
field_type_base = PROTO_TO_GRAPHENE_TYPE[field.type]
|
|
if field.label == FieldDescriptor.LABEL_REPEATED:
|
|
return f"graphene.List({field_type_base})"
|
|
else:
|
|
return f"{field_type_base}()"
|
|
|
|
|
|
def proto_method_to_graphql_operation(method):
|
|
method_name = get_method_name(method)
|
|
input_descriptor = method.input_type
|
|
output_descriptor = method.output_type
|
|
input_class_name = get_descriptor_full_pascal_name(input_descriptor) + "Input"
|
|
out_put_class_name = get_descriptor_full_pascal_name(output_descriptor)
|
|
field_def = f"graphene.Field({out_put_class_name}, input={input_class_name}())"
|
|
return f"\n{INDENT}{method_name} = {field_def}"
|
|
|
|
|
|
def generate_resolver_function(method):
|
|
full_method_name = get_method_name(method)
|
|
snake_case_method_name = camel_to_snake(method.name)
|
|
pascal_case_method_name = snake_to_pascal(snake_case_method_name)
|
|
pb2_file_name = method_descriptor_to_generated_pb2_file_name(method)
|
|
|
|
function_builder = ""
|
|
function_builder += f"\n{INDENT}def resolve_{full_method_name}(self, info, input):"
|
|
function_builder += f"\n{INDENT2}input_dict = vars(input)"
|
|
function_builder += (
|
|
f"\n{INDENT2}request_message = mlflow.protos.{pb2_file_name}.{pascal_case_method_name}()"
|
|
)
|
|
function_builder += f"\n{INDENT2}parse_dict(input_dict, request_message)"
|
|
function_builder += (
|
|
f"\n{INDENT2}return mlflow.server.handlers.{snake_case_method_name}_impl(request_message)"
|
|
)
|
|
function_builder += "\n"
|
|
return function_builder
|