H-3225: Allow querying of closed data schemas (#5527)

Cargo.toml

@@ -209,6 +209,7 @@ tokio-stream = { version = "=0.1.16", default-features = false }
 tokio-test = { version = "=0.4.4", default-features = false }
 tower = { version = "=0.5.1", default-features = false }
 tower-test = { version = "=0.4.0", default-features = false }
+trait-variant = { version = "=0.1.2", default-features = false }
 tracing = { version = "=0.1.40", default-features = false }
 tracing-error = { version = "=0.2.0", default-features = false }
 tracing-flame = { version = "=0.2.0", default-features = false }

apps/hash-graph/libs/graph/src/store/postgres/knowledge/entity/mod.rs

@@ -31,7 +31,7 @@ use graph_types::{
             PropertyWithMetadataValue, visitor::EntityVisitor as _,
         },
     },
-    ontology::{DataTypeProvider, OntologyTypeProvider},
+    ontology::{DataTypeLookup, OntologyTypeProvider},
     owned_by_id::OwnedById,
 };
 use hash_graph_store::{
@@ -60,7 +60,8 @@ use temporal_versioning::{
 use tokio_postgres::{GenericClient as _, Row, error::SqlState};
 use type_system::{
     schema::{
-        ClosedEntityType, ClosedMultiEntityType, EntityTypeUuid, InheritanceDepth, OntologyTypeUuid,
+        ClosedEntityType, ClosedMultiEntityType, DataTypeReference, EntityTypeUuid,
+        InheritanceDepth, OntologyTypeUuid,
     },
     url::VersionedUrl,
 };
@@ -333,7 +334,7 @@ where
         Ok(())
     }
 
-    async fn convert_entity_properties<P: DataTypeProvider + Sync>(
+    async fn convert_entity_properties<P: DataTypeLookup + Sync>(
         &self,
         provider: &P,
         entity: &mut PropertyWithMetadata,
@@ -353,7 +354,10 @@ where
         };
 
         let Ok(conversions) = provider
-            .find_conversion(source_data_type_id, target_data_type_id)
+            .find_conversion(
+                <&DataTypeReference>::from(&*source_data_type_id),
+                <&DataTypeReference>::from(target_data_type_id),
+            )
             .await
         else {
             // If no conversion is found, we can ignore the property.
@@ -375,7 +379,7 @@ where
         metadata.data_type_id = Some(target_data_type_id.clone());
     }
 
-    async fn convert_entity<P: DataTypeProvider + Sync>(
+    async fn convert_entity<P: DataTypeLookup + Sync>(
         &self,
         provider: &P,
         entity: &mut Entity,

apps/hash-graph/libs/graph/src/store/validation.rs

@@ -14,7 +14,7 @@ use graph_types::{
     account::AccountId,
     knowledge::entity::{Entity, EntityId},
     ontology::{
-        DataTypeProvider, DataTypeWithMetadata, EntityTypeProvider, EntityTypeWithMetadata,
+        DataTypeLookup, DataTypeWithMetadata, EntityTypeProvider, EntityTypeWithMetadata,
         OntologyTypeProvider, PropertyTypeProvider, PropertyTypeWithMetadata,
     },
 };
@@ -28,9 +28,10 @@ use hash_graph_store::{
 use tokio::sync::RwLock;
 use tokio_postgres::GenericClient as _;
 use type_system::{
+    Valid,
     schema::{
-        ClosedEntityType, ConversionDefinition, ConversionExpression, DataTypeUuid, EntityTypeUuid,
-        PropertyType, PropertyTypeUuid,
+        ClosedDataType, ClosedEntityType, ConversionDefinition, ConversionExpression,
+        DataTypeReference, DataTypeUuid, EntityTypeUuid, PropertyType, PropertyTypeUuid,
     },
     url::{BaseUrl, VersionedUrl},
 };
@@ -123,6 +124,7 @@ where
 #[derive(Debug, Default)]
 pub struct StoreCache {
     data_types: CacheHashMap<DataTypeUuid, DataTypeWithMetadata>,
+    closed_data_types: CacheHashMap<DataTypeUuid, ClosedDataType>,
     property_types: CacheHashMap<PropertyTypeUuid, PropertyType>,
     entity_types: CacheHashMap<EntityTypeUuid, ClosedEntityType>,
     entities: CacheHashMap<EntityId, Entity>,
@@ -161,33 +163,32 @@ where
     }
 }
 
-impl<C, A> OntologyTypeProvider<DataTypeWithMetadata> for StoreProvider<'_, PostgresStore<C, A>>
+impl<C, A> DataTypeLookup for StoreProvider<'_, PostgresStore<C, A>>
 where
     C: AsClient,
     A: AuthorizationApi,
 {
-    type Value = Arc<DataTypeWithMetadata>;
+    type ClosedDataType = Arc<ClosedDataType>;
+    type DataTypeWithMetadata = Arc<DataTypeWithMetadata>;
+    type Error = QueryError;
 
-    #[expect(refining_impl_trait)]
-    async fn provide_type(
+    async fn lookup_data_type_by_uuid(
         &self,
-        type_id: &VersionedUrl,
+        data_type_uuid: DataTypeUuid,
     ) -> Result<Arc<DataTypeWithMetadata>, Report<QueryError>> {
-        let data_type_id = DataTypeUuid::from_url(type_id);
-
-        if let Some(cached) = self.cache.data_types.get(&data_type_id).await {
+        if let Some(cached) = self.cache.data_types.get(&data_type_uuid).await {
             return cached;
         }
 
-        if let Err(error) = self.authorize_data_type(data_type_id).await {
-            self.cache.data_types.deny(data_type_id).await;
+        if let Err(error) = self.authorize_data_type(data_type_uuid).await {
+            self.cache.data_types.deny(data_type_uuid).await;
             return Err(error);
         }
 
         let schema = self
             .store
             .read_one(
-                &Filter::<DataTypeWithMetadata>::for_versioned_url(type_id),
+                &Filter::for_data_type_uuid(data_type_uuid),
                 Some(
                     &QueryTemporalAxesUnresolved::DecisionTime {
                         pinned: PinnedTemporalAxisUnresolved::new(None),
@@ -199,25 +200,51 @@ where
             )
             .await?;
 
-        let schema = self.cache.data_types.grant(data_type_id, schema).await;
+        let schema = self.cache.data_types.grant(data_type_uuid, schema).await;
+
+        Ok(schema)
+    }
+
+    async fn lookup_closed_data_type_by_uuid(
+        &self,
+        data_type_uuid: DataTypeUuid,
+    ) -> Result<Arc<ClosedDataType>, Report<QueryError>> {
+        if let Some(cached) = self.cache.closed_data_types.get(&data_type_uuid).await {
+            return cached;
+        }
+
+        if let Err(error) = self.authorize_data_type(data_type_uuid).await {
+            self.cache.closed_data_types.deny(data_type_uuid).await;
+            return Err(error);
+        }
+
+        let schema: Valid<ClosedDataType> = self
+            .store
+            .as_client()
+            .query_one(
+                "SELECT closed_schema FROM data_types WHERE ontology_id = $1",
+                &[&data_type_uuid],
+            )
+            .await
+            .change_context(QueryError)?
+            .get(0);
+
+        let schema = self
+            .cache
+            .closed_data_types
+            .grant(data_type_uuid, schema.into_inner())
+            .await;
 
         Ok(schema)
     }
-}
 
-impl<C, A> DataTypeProvider for StoreProvider<'_, PostgresStore<C, A>>
-where
-    C: AsClient,
-    A: AuthorizationApi,
-{
-    #[expect(refining_impl_trait)]
     async fn is_parent_of(
         &self,
-        child: &VersionedUrl,
+        child: &DataTypeReference,
         parent: &BaseUrl,
     ) -> Result<bool, Report<QueryError>> {
         let client = self.store.as_client().client();
-        let child = DataTypeUuid::from_url(child);
+        let child = DataTypeUuid::from_url(&child.url);
 
         Ok(client
             .query_one(
@@ -237,16 +264,20 @@ where
             .get(0))
     }
 
-    #[expect(refining_impl_trait)]
     async fn find_conversion(
         &self,
-        source_data_type_id: &VersionedUrl,
-        target_data_type_id: &VersionedUrl,
+        source: &DataTypeReference,
+        target: &DataTypeReference,
     ) -> Result<impl Borrow<Vec<ConversionExpression>>, Report<QueryError>> {
-        let source = DataTypeUuid::from_url(source_data_type_id);
-        let target = DataTypeUuid::from_url(target_data_type_id);
+        let source_uuid = DataTypeUuid::from_url(&source.url);
+        let target_uuid = DataTypeUuid::from_url(&target.url);
 
-        if let Some(cached) = self.cache.conversions.get(&(source, target)).await {
+        if let Some(cached) = self
+            .cache
+            .conversions
+            .get(&(source_uuid, target_uuid))
+            .await
+        {
             return cached;
         }
 
@@ -274,18 +305,18 @@ where
                        AND target_data_type_base_url = $3
                 ;",
                 &[
-                    &source,
-                    &target,
-                    &source_data_type_id.base_url,
-                    &target_data_type_id.base_url,
+                    &source_uuid,
+                    &target_uuid,
+                    &source.url.base_url,
+                    &target.url.base_url,
                 ],
             )
             .await
             .change_context(QueryError)
             .attach_printable_lazy(|| {
                 format!(
-                    "Found none or more than one conversions between `{source_data_type_id}` and \
-                     `{target_data_type_id}`"
+                    "Found none or more than one conversions between `{}` and `{}`",
+                    source.url, target.url
                 )
             })?
             .get::<_, Vec<ConversionDefinition>>(0)
@@ -296,7 +327,7 @@ where
         Ok(self
             .cache
             .conversions
-            .grant((source, target), expression)
+            .grant((source_uuid, target_uuid), expression)
             .await)
     }
 }