wehub-resource-sync
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import os
|
|
|
|
from autogeneration_utils import AUTOGENERATED_SCHEMA, AUTOGENERATED_SDL_SCHEMA
|
|
from parsing_utils import process_method
|
|
from schema_autogeneration import generate_schema
|
|
|
|
from mlflow.protos import model_registry_pb2, service_pb2
|
|
from mlflow.server.graphql.graphql_schema_extensions import schema
|
|
|
|
# Add proto descriptors to onboard RPCs to graphql.
|
|
ONBOARDED_DESCRIPTORS = [service_pb2.DESCRIPTOR, model_registry_pb2.DESCRIPTOR]
|
|
|
|
|
|
class GenerateSchemaState:
|
|
def __init__(self):
|
|
self.queries = set() # method_descriptor
|
|
self.mutations = set() # method_descriptor
|
|
self.inputs = [] # field_descriptor
|
|
self.outputs = set() # field_descriptor
|
|
self.types = [] # field_descriptor
|
|
self.enums = set() # enum_descriptor
|
|
self.method_names = set() # package_name_method_name
|
|
|
|
|
|
# Entry point for generating the GraphQL schema.
|
|
def generate_code():
|
|
state = GenerateSchemaState()
|
|
for file_descriptor in ONBOARDED_DESCRIPTORS:
|
|
for service_name, service_descriptor in file_descriptor.services_by_name.items():
|
|
for method_name, method_descriptor in service_descriptor.methods_by_name.items():
|
|
process_method(method_descriptor, state)
|
|
|
|
generated_schema = generate_schema(state)
|
|
|
|
os.makedirs(os.path.dirname(AUTOGENERATED_SCHEMA), exist_ok=True)
|
|
|
|
with open(AUTOGENERATED_SCHEMA, "w") as file:
|
|
file.write(generated_schema)
|
|
|
|
# Generate the sdl schema for typescript type generation.
|
|
sdl_schema = str(schema)
|
|
sdl_schema = f"""# GENERATED FILE. PLEASE DON'T MODIFY.
|
|
# Run uv run ./dev/proto_to_graphql/code_generator.py to regenerate.
|
|
|
|
{sdl_schema}
|
|
|
|
"""
|
|
|
|
with open(AUTOGENERATED_SDL_SCHEMA, "w") as f:
|
|
f.write(sdl_schema)
|
|
|
|
|
|
def main():
|
|
generate_code()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|