SchemaRegistry::getTypeSystemDefinitionRefByUri

Summary: Enables accessing bundled schema data as TypeSystem nodes by uri, for definitions opted in using the `any` cpp2 compiler option, or statically for any definitions.

Reviewed By: pranavtbhat

Differential Revision: D75178792

fbshipit-source-id: 8f36fb30ef5def1a52c8bc65233fd154ff01f939

Author: iahs

third-party/thrift/src/thrift/lib/cpp2/runtime/SchemaRegistry.cpp

@@ -102,6 +102,16 @@ SchemaRegistry::getSyntaxGraphDefinitionNodeByUri(
       uri, type::ProgramId{data->programId}, data->data);
 }
 
+std::optional<type_system::DefinitionRef>
+SchemaRegistry::getTypeSystemDefinitionRefByUri(
+    const std::string_view uri) const {
+  if (const auto* sgDef = getSyntaxGraphDefinitionNodeByUri(uri)) {
+    std::lock_guard<folly::SharedMutex> wlock(typeSystemMutex_);
+    return syntaxGraph_->asTypeSystemDefinitionRef(*sgDef);
+  }
+  return std::nullopt;
+}
+
 } // namespace apache::thrift
 
 #endif

third-party/thrift/src/thrift/lib/cpp2/runtime/SchemaRegistry.h

@@ -21,8 +21,10 @@
 #ifdef THRIFT_SCHEMA_AVAILABLE
 
 #include <unordered_set>
+#include <folly/SharedMutex.h>
 #include <folly/synchronization/RelaxedAtomic.h>
 #include <thrift/lib/cpp/util/EnumUtils.h>
+#include <thrift/lib/cpp2/dynamic/TypeSystem.h>
 #include <thrift/lib/cpp2/runtime/BaseSchemaRegistry.h>
 #include <thrift/lib/cpp2/schema/SyntaxGraph.h>
 #include <thrift/lib/cpp2/schema/detail/SchemaBackedResolver.h>
@@ -31,32 +33,49 @@
 namespace apache::thrift {
 namespace detail {
 template <typename T>
-struct SyntaxGraphNodeTag {
+struct NodeTag {
   using type = T;
 };
 
 template <typename T>
 constexpr auto getSyntaxGraphNodeTypeFor() {
   if constexpr (is_thrift_service_tag_v<T>) {
-    return SyntaxGraphNodeTag<syntax_graph::ServiceNode>{};
+    return NodeTag<syntax_graph::ServiceNode>{};
   } else if constexpr (is_thrift_struct_v<T>) {
-    return SyntaxGraphNodeTag<syntax_graph::StructNode>{};
+    return NodeTag<syntax_graph::StructNode>{};
   } else if constexpr (is_thrift_union_v<T>) {
-    return SyntaxGraphNodeTag<syntax_graph::UnionNode>{};
+    return NodeTag<syntax_graph::UnionNode>{};
   } else if constexpr (is_thrift_exception_v<T>) {
-    return SyntaxGraphNodeTag<syntax_graph::ExceptionNode>{};
+    return NodeTag<syntax_graph::ExceptionNode>{};
   } else if constexpr (util::is_thrift_enum_v<T>) {
-    return SyntaxGraphNodeTag<syntax_graph::EnumNode>{};
+    return NodeTag<syntax_graph::EnumNode>{};
   } else {
     static_assert(folly::always_false<T>, "Unsupported Thrift type");
   }
   // It's unclear how to include typedefs and constants here.
   // TODO: interactions
 }
 
+template <typename T>
+constexpr auto getTypeSystemNodeTypeFor() {
+  if constexpr (is_thrift_struct_v<T>) {
+    return NodeTag<type_system::StructNode>{};
+  } else if constexpr (is_thrift_union_v<T>) {
+    return NodeTag<type_system::UnionNode>{};
+  } else if constexpr (util::is_thrift_enum_v<T>) {
+    return NodeTag<type_system::EnumNode>{};
+  } else {
+    static_assert(folly::always_false<T>, "Unsupported Thrift type");
+  }
+}
+
 template <typename T>
 using SyntaxGraphNodeTypeFor =
     typename decltype(getSyntaxGraphNodeTypeFor<T>())::type;
+
+template <typename T>
+using TypeSystemNodeTypeFor =
+    typename decltype(getTypeSystemNodeTypeFor<T>())::type;
 } // namespace detail
 namespace test {
 struct SchemaTest;
@@ -88,6 +107,36 @@ class SchemaRegistry {
         .template as<detail::SyntaxGraphNodeTypeFor<T>>();
   }
 
+  /**
+   * Gets TypeSystem node for given URI, or nullopt if not found.
+   * Only types defined in a file with the `any` cpp2 compiler option enabled
+   * can be found using this method.
+   */
+  std::optional<type_system::DefinitionRef> getTypeSystemDefinitionRefByUri(
+      const std::string_view uri) const;
+
+  /**
+   * Gets TypeSystem node for given definition, or throws `std::out_of_range` if
+   * not present in schema.
+   */
+  template <typename T>
+  type_system::DefinitionRef getTypeSystemDefinitionRef() const {
+    const syntax_graph::DefinitionNode& sgNode = getDefinitionNode<T>();
+    std::lock_guard<folly::SharedMutex> wlock(typeSystemMutex_);
+    return syntaxGraph_->asTypeSystemDefinitionRef(sgNode);
+  }
+
+  /**
+   * Gets TypeSystem node for given definition, or throws `std::out_of_range` if
+   * not present in schema. Returns most-derived type (e.g. StructNode) that
+   * matches the template parameter.
+   */
+  template <typename T>
+  const detail::TypeSystemNodeTypeFor<T>& getTypeSystemNode() const {
+    return getTypeSystemDefinitionRef<T>()
+        .template asType<detail::TypeSystemNodeTypeFor<T>>();
+  }
+
   explicit SchemaRegistry(BaseSchemaRegistry& base);
   ~SchemaRegistry();
 
@@ -111,6 +160,9 @@ class SchemaRegistry {
 
   std::unique_ptr<syntax_graph::SyntaxGraph> syntaxGraph_;
   syntax_graph::detail::IncrementalResolver* resolver_;
+  // Synchronizes access to syntaxGraph_->asTypeSystem* methods, which are not
+  // thread-safe.
+  mutable folly::SharedMutex typeSystemMutex_;
 
   friend struct test::SchemaTest;
 };

third-party/thrift/src/thrift/test/SchemaTest.cpp

@@ -217,4 +217,29 @@ TEST_F(SchemaTest, getSyntaxGraphDefinitionNodeByUri) {
   EXPECT_FALSE(getSyntaxGraphDefinitionNodeByUri(uriWithoutAny));
 }
 
+TEST_F(SchemaTest, getTypeSystemDefinitionRef) {
+  auto& registry = SchemaRegistry::get();
+
+  auto uri = "facebook.com/thrift/test/schema/Empty";
+
+  auto dynamicDefinitionRef = registry.getTypeSystemDefinitionRefByUri(uri);
+  EXPECT_TRUE(dynamicDefinitionRef);
+
+  const type_system::StructNode& dynamicStructNode =
+      dynamicDefinitionRef->asStruct();
+  EXPECT_EQ(uri, dynamicStructNode.uri());
+
+  auto staticDefinitionRef =
+      registry
+          .getTypeSystemDefinitionRef<facebook::thrift::test::schema::Empty>();
+  const type_system::StructNode& staticStructNode =
+      staticDefinitionRef.asStruct();
+
+  const type_system::StructNode& directStructNode =
+      registry.getTypeSystemNode<facebook::thrift::test::schema::Empty>();
+
+  EXPECT_EQ(&dynamicStructNode, &staticStructNode);
+  EXPECT_EQ(&dynamicStructNode, &directStructNode);
+}
+
 } // namespace apache::thrift::test