chore: import upstream snapshot with attribution

dev/proto_to_graphql/code_generator.py

@@ -0,0 +1,58 @@
+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()