wehub-resource-sync
68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
from autogeneration_utils import get_method_name
|
|
from google.protobuf.descriptor import FieldDescriptor
|
|
|
|
from mlflow.protos import databricks_pb2
|
|
|
|
|
|
def get_method_type(method_descriptor):
|
|
return method_descriptor.GetOptions().Extensions[databricks_pb2.rpc].endpoints[0].method
|
|
|
|
|
|
def process_method(method_descriptor, state):
|
|
"""
|
|
Given a method descriptor, add information being referenced into the GenerateSchemaState.
|
|
"""
|
|
if not method_descriptor.GetOptions().HasExtension(databricks_pb2.graphql):
|
|
return
|
|
rpcOptions = method_descriptor.GetOptions().Extensions[databricks_pb2.rpc]
|
|
# Only add those methods that are not internal.
|
|
if rpcOptions.visibility != databricks_pb2.INTERNAL:
|
|
name = get_method_name(method_descriptor)
|
|
if name in state.method_names:
|
|
return
|
|
state.method_names.add(name)
|
|
request_method = get_method_type(method_descriptor)
|
|
if request_method == "GET":
|
|
state.queries.add(method_descriptor)
|
|
else:
|
|
state.mutations.add(method_descriptor)
|
|
state.outputs.add(method_descriptor.output_type)
|
|
populate_message_types(method_descriptor.input_type, state, True, set())
|
|
populate_message_types(method_descriptor.output_type, state, False, set())
|
|
|
|
|
|
def populate_message_types(field_descriptor, state, is_input, visited):
|
|
"""
|
|
Given a field descriptor, recursively walk through the referenced message types and add
|
|
information being referenced into the GenerateSchemaState.
|
|
"""
|
|
if field_descriptor in visited:
|
|
# Break the loop for recursive types.
|
|
return
|
|
visited.add(field_descriptor)
|
|
if is_input:
|
|
add_message_descriptor_to_list(field_descriptor, state.inputs)
|
|
else:
|
|
add_message_descriptor_to_list(field_descriptor, state.types)
|
|
|
|
for sub_field in field_descriptor.fields:
|
|
type = sub_field.type
|
|
if type in (FieldDescriptor.TYPE_MESSAGE, FieldDescriptor.TYPE_GROUP):
|
|
populate_message_types(sub_field.message_type, state, is_input, visited)
|
|
elif type == FieldDescriptor.TYPE_ENUM:
|
|
state.enums.add(sub_field.enum_type)
|
|
else:
|
|
continue
|
|
|
|
|
|
def add_message_descriptor_to_list(descriptor, target_list):
|
|
# Always put the referenced message at the beginning, so that when generating the schema,
|
|
# the ordering can be maintained in a way that correspond to the reference graph.
|
|
# list.remove() and insert(0) are not optimal in terms of efficiency but are fine because
|
|
# the amount of data is very small here.
|
|
if descriptor not in target_list:
|
|
target_list.insert(0, descriptor)
|
|
else:
|
|
target_list.remove(descriptor)
|
|
target_list.insert(0, descriptor)
|