H-3543: Add interface to read closed multi-entity type schemas

apps/hash-graph/libs/api/src/rest/entity.rs

@@ -25,10 +25,11 @@ use graph::store::{
     NullOrdering, Ordering, StorePool,
     error::{EntityDoesNotExist, RaceConditionOnUpdate},
     knowledge::{
-        CountEntitiesParams, CreateEntityRequest, DiffEntityParams, DiffEntityResult,
-        EntityQueryCursor, EntityQuerySorting, EntityQuerySortingRecord, EntityStore as _,
-        EntityValidationType, GetEntitiesParams, GetEntitiesResponse, GetEntitySubgraphParams,
-        PatchEntityParams, QueryConversion, UpdateEntityEmbeddingsParams, ValidateEntityParams,
+        ClosedMultiEntityTypeMap, CountEntitiesParams, CreateEntityRequest, DiffEntityParams,
+        DiffEntityResult, EntityQueryCursor, EntityQuerySorting, EntityQuerySortingRecord,
+        EntityStore as _, EntityValidationType, GetEntitiesParams, GetEntitiesResponse,
+        GetEntitySubgraphParams, PatchEntityParams, QueryConversion, UpdateEntityEmbeddingsParams,
+        ValidateEntityParams,
     },
 };
 use graph_types::{
@@ -132,6 +133,7 @@ use crate::rest::{
             EntityQuerySortingToken,
             GetEntitiesResponse,
             GetEntitySubgraphResponse,
+            ClosedMultiEntityTypeMap,
             QueryConversion,
 
             Entity,
@@ -538,6 +540,8 @@ struct GetEntitiesRequest<'q, 's, 'p> {
     #[serde(default)]
     include_count: bool,
     #[serde(default)]
+    include_closed_multi_entity_types: bool,
+    #[serde(default)]
     include_web_ids: bool,
     #[serde(default)]
     include_created_by_ids: bool,
@@ -609,6 +613,7 @@ where
             conversions: request.conversions,
             include_drafts: request.include_drafts,
             include_count: request.include_count,
+            include_closed_multi_entity_types: request.include_closed_multi_entity_types,
             temporal_axes: request.temporal_axes,
             include_web_ids: request.include_web_ids,
             include_created_by_ids: request.include_created_by_ids,
@@ -621,6 +626,7 @@ where
                 entities: response.entities,
                 cursor: response.cursor.map(EntityQueryCursor::into_owned),
                 count: response.count,
+                closed_multi_entity_types: response.closed_multi_entity_types,
                 web_ids: response.web_ids,
                 created_by_ids: response.created_by_ids,
                 edition_created_by_ids: response.edition_created_by_ids,
@@ -652,6 +658,8 @@ struct GetEntitySubgraphRequest<'q, 's, 'p> {
     #[serde(default)]
     include_count: bool,
     #[serde(default)]
+    include_closed_multi_entity_types: bool,
+    #[serde(default)]
     include_web_ids: bool,
     #[serde(default)]
     include_created_by_ids: bool,
@@ -668,6 +676,8 @@ struct GetEntitySubgraphResponse<'r> {
     #[serde(borrow)]
     cursor: Option<EntityQueryCursor<'r>>,
     count: Option<usize>,
+    #[serde(skip_serializing_if = "HashMap::is_empty")]
+    closed_multi_entity_types: HashMap<VersionedUrl, ClosedMultiEntityTypeMap>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[schema(nullable = false)]
     web_ids: Option<HashMap<OwnedById, usize>>,
@@ -745,6 +755,7 @@ where
             graph_resolve_depths: request.graph_resolve_depths,
             include_drafts: request.include_drafts,
             include_count: request.include_count,
+            include_closed_multi_entity_types: request.include_closed_multi_entity_types,
             temporal_axes: request.temporal_axes,
             include_web_ids: request.include_web_ids,
             include_created_by_ids: request.include_created_by_ids,
@@ -757,6 +768,7 @@ where
                 subgraph: response.subgraph.into(),
                 cursor: response.cursor.map(EntityQueryCursor::into_owned),
                 count: response.count,
+                closed_multi_entity_types: response.closed_multi_entity_types,
                 web_ids: response.web_ids,
                 created_by_ids: response.created_by_ids,
                 edition_created_by_ids: response.edition_created_by_ids,

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

@@ -225,19 +225,32 @@ pub struct GetEntitiesParams<'a> {
     pub limit: Option<usize>,
     pub include_drafts: bool,
     pub include_count: bool,
+    pub include_closed_multi_entity_types: bool,
     pub include_web_ids: bool,
     pub include_created_by_ids: bool,
     pub include_edition_created_by_ids: bool,
     pub include_type_ids: bool,
 }
 
+#[derive(Debug, Serialize)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde(rename_all = "camelCase")]
+pub struct ClosedMultiEntityTypeMap {
+    #[cfg_attr(feature = "utoipa", schema(value_type = VAR_CLOSED_MULTI_ENTITY_TYPE))]
+    pub schema: ClosedMultiEntityType,
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub inner: HashMap<VersionedUrl, Self>,
+}
+
 #[derive(Debug, Serialize)]
 #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
 #[serde(rename_all = "camelCase")]
 pub struct GetEntitiesResponse<'r> {
     pub entities: Vec<Entity>,
     pub cursor: Option<EntityQueryCursor<'r>>,
     pub count: Option<usize>,
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub closed_multi_entity_types: HashMap<VersionedUrl, ClosedMultiEntityTypeMap>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[cfg_attr(feature = "utoipa", schema(nullable = false))]
     pub web_ids: Option<HashMap<OwnedById, usize>>,
@@ -263,6 +276,7 @@ pub struct GetEntitySubgraphParams<'a> {
     pub conversions: Vec<QueryConversion<'a>>,
     pub include_drafts: bool,
     pub include_count: bool,
+    pub include_closed_multi_entity_types: bool,
     pub include_web_ids: bool,
     pub include_created_by_ids: bool,
     pub include_edition_created_by_ids: bool,
@@ -274,6 +288,7 @@ pub struct GetEntitySubgraphResponse<'r> {
     pub subgraph: Subgraph,
     pub cursor: Option<EntityQueryCursor<'r>>,
     pub count: Option<usize>,
+    pub closed_multi_entity_types: HashMap<VersionedUrl, ClosedMultiEntityTypeMap>,
     pub web_ids: Option<HashMap<OwnedById, usize>>,
     pub created_by_ids: Option<HashMap<CreatedById, usize>>,
     pub edition_created_by_ids: Option<HashMap<EditionCreatedById, usize>>,

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

@@ -1,6 +1,6 @@
 mod query;
 mod read;
-use alloc::borrow::Cow;
+use alloc::{borrow::Cow, collections::BTreeSet};
 use core::{borrow::Borrow as _, iter::once, mem};
 use std::collections::{HashMap, HashSet};
 
@@ -72,9 +72,9 @@ use crate::store::{
     crud::{QueryResult as _, Read, ReadPaginated, Sorting as _},
     error::{DeletionError, EntityDoesNotExist, RaceConditionOnUpdate},
     knowledge::{
-        CountEntitiesParams, CreateEntityParams, EntityQuerySorting, EntityStore,
-        EntityValidationType, GetEntitiesParams, GetEntitiesResponse, GetEntitySubgraphParams,
-        GetEntitySubgraphResponse, PatchEntityParams, QueryConversion,
+        ClosedMultiEntityTypeMap, CountEntitiesParams, CreateEntityParams, EntityQuerySorting,
+        EntityStore, EntityValidationType, GetEntitiesParams, GetEntitiesResponse,
+        GetEntitySubgraphParams, GetEntitySubgraphResponse, PatchEntityParams, QueryConversion,
         UpdateEntityEmbeddingsParams, ValidateEntityError, ValidateEntityParams,
     },
     postgres::{
@@ -99,6 +99,7 @@ struct GetEntitiesImplParams<'a> {
     limit: Option<usize>,
     include_drafts: bool,
     include_count: bool,
+    include_closed_multi_entity_types: bool,
     include_web_ids: bool,
     include_created_by_ids: bool,
     include_edition_created_by_ids: bool,
@@ -403,6 +404,92 @@ where
         Ok(())
     }
 
+    // #[tracing::instrument(level = "info", skip(self, entities))]
+    async fn resolve_closed_multi_entity_types(
+        &self,
+        entities: impl IntoIterator<Item = &Entity> + Send,
+    ) -> Result<HashMap<VersionedUrl, ClosedMultiEntityTypeMap>, Report<QueryError>> {
+        let mut entity_type_ids_to_resolve = HashSet::<&VersionedUrl>::new();
+        let all_multi_entity_type_ids = entities
+            .into_iter()
+            .map(|entity| {
+                entity_type_ids_to_resolve.extend(entity.metadata.entity_type_ids.iter());
+                entity
+                    .metadata
+                    .entity_type_ids
+                    .iter()
+                    .collect::<BTreeSet<_>>()
+            })
+            .collect::<Vec<_>>();
+
+        let (entity_type_references, entity_type_uuids): (Vec<_>, Vec<_>) =
+            entity_type_ids_to_resolve
+                .into_iter()
+                .map(|entity_type_id| (EntityTypeUuid::from_url(entity_type_id), entity_type_id))
+                .unzip();
+
+        let closed_types = entity_type_uuids
+            .into_iter()
+            .zip(
+                self.get_closed_entity_types(&entity_type_references)
+                    .await?,
+            )
+            .collect::<HashMap<_, _>>();
+
+        let mut resolved_entity_types = HashMap::<VersionedUrl, ClosedMultiEntityTypeMap>::new();
+        for entity_multi_type_ids in all_multi_entity_type_ids {
+            let mut entity_type_id_iter = entity_multi_type_ids.iter();
+            let Some(first_entity_type_id) = entity_type_id_iter.next() else {
+                continue;
+            };
+            let (_, ref mut map) = resolved_entity_types
+                .raw_entry_mut()
+                .from_key(*first_entity_type_id)
+                .or_insert_with(|| {
+                    ((*first_entity_type_id).clone(), ClosedMultiEntityTypeMap {
+                        schema: ClosedMultiEntityType::from_closed_schema(
+                            closed_types
+                                .get(*first_entity_type_id)
+                                .expect(
+                                    "The entity type was already resolved, so it should be \
+                                     present in the closed types",
+                                )
+                                .clone(),
+                        ),
+                        inner: HashMap::new(),
+                    })
+                });
+
+            for entity_type_id in entity_type_id_iter {
+                let (_, new_map) = map
+                    .inner
+                    .raw_entry_mut()
+                    .from_key(*entity_type_id)
+                    .or_insert_with(|| {
+                        let mut closed_parent = map.schema.clone();
+                        closed_parent
+                            .add_closed_entity_type(
+                                closed_types
+                                    .get(*entity_type_id)
+                                    .expect(
+                                        "The entity type was already resolved, so it should be \
+                                         present in the closed types",
+                                    )
+                                    .clone(),
+                            )
+                            .expect("The entity type was constructed before so it has to be valid");
+                        ((*entity_type_id).clone(), ClosedMultiEntityTypeMap {
+                            schema: closed_parent,
+                            inner: HashMap::new(),
+                        })
+                    });
+                *map = new_map;
+            }
+        }
+
+        Ok(resolved_entity_types)
+    }
+
     #[tracing::instrument(level = "info", skip(self, params))]
     async fn get_entities_impl(
         &self,
@@ -557,14 +644,9 @@ where
                 )
             };
 
-            root_entities.extend(
-                entities
-                    .into_iter()
-                    .filter(|(entity, _)| {
-                        permitted_ids.contains(&entity.metadata.record_id.entity_id.entity_uuid)
-                    })
-                    .take(params.limit.unwrap_or(usize::MAX) - root_entities.len()),
-            );
+            root_entities.extend(entities.into_iter().filter(|(entity, _)| {
+                permitted_ids.contains(&entity.metadata.record_id.entity_id.entity_uuid)
+            }));
 
             if let Some(limit) = params.limit {
                 if num_returned_entities < limit {
@@ -589,6 +671,14 @@ where
 
         Ok((
             GetEntitiesResponse {
+                closed_multi_entity_types: if params.include_closed_multi_entity_types {
+                    self.resolve_closed_multi_entity_types(
+                        root_entities.iter().map(|(entity, _)| entity),
+                    )
+                    .await?
+                } else {
+                    HashMap::new()
+                },
                 entities: root_entities
                     .into_iter()
                     .map(|(entity, _)| entity)
@@ -1111,6 +1201,7 @@ where
                     limit: params.limit,
                     include_drafts: params.include_drafts,
                     include_count: params.include_count,
+                    include_closed_multi_entity_types: params.include_closed_multi_entity_types,
                     include_web_ids: params.include_web_ids,
                     include_created_by_ids: params.include_created_by_ids,
                     include_edition_created_by_ids: params.include_edition_created_by_ids,
@@ -1163,6 +1254,7 @@ where
                 entities: root_entities,
                 cursor,
                 count,
+                closed_multi_entity_types,
                 web_ids,
                 created_by_ids,
                 edition_created_by_ids,
@@ -1178,6 +1270,7 @@ where
                     limit: params.limit,
                     include_drafts: params.include_drafts,
                     include_count: false,
+                    include_closed_multi_entity_types: params.include_closed_multi_entity_types,
                     include_web_ids: params.include_web_ids,
                     include_created_by_ids: params.include_created_by_ids,
                     include_edition_created_by_ids: params.include_edition_created_by_ids,
@@ -1251,6 +1344,7 @@ where
             subgraph,
             cursor,
             count,
+            closed_multi_entity_types,
             web_ids,
             created_by_ids,
             edition_created_by_ids,