feat: Also fetch domains and data products

Author: sbernauer

rust/resource-info-fetcher/src/backend/data_hub/graphql.rs

@@ -8,26 +8,42 @@ use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
 
-use crate::backend::data_hub::{DataHubResourceInfoResponse, Group, Owners, Tag, Urn, User};
+use crate::backend::data_hub::{
+    DataHubResourceInfoResponse, DataProduct, Domain, Group, Owners, Tag, Urn, User,
+};
 
 /// A single query covering every entity kind we build URNs for. We use the generic `entity(urn:)`
 /// resolver plus per-type inline fragments, because a request can target a dataset (Trino table or
 /// Kafka topic), a container (Trino catalog or schema), a chart or a dashboard.
-const RESOURCE_INFO_QUERY: &str = r"
+const RESOURCE_INFO_QUERY: &str = r#"
 query ResourceInfo($urn: String!) {
   entity(urn: $urn) {
     ...ResourceInfo
   }
 }
 fragment ResourceInfo on Entity {
-  ... on Dataset   { tags { ...Tags } ownership { ...Owners } }
-  ... on Container { tags { ...Tags } ownership { ...Owners } }
-  ... on Chart     { tags { ...Tags } ownership { ...Owners } }
-  ... on Dashboard { tags { ...Tags } ownership { ...Owners } }
+  # DataHub has no direct "dataProduct" field on assets; membership is a graph edge that points from
+  # the data product to its assets. From the asset's side it is therefore an INCOMING relationship.
+  dataProducts: relationships(input: {types: ["DataProductContains"], direction: INCOMING, count: 10}) {
+    relationships { ...DataProduct }
+  }
+  ... on Dataset   { tags { ...Tags } ownership { ...Owners } domain { ...Domain } }
+  ... on Container { tags { ...Tags } ownership { ...Owners } domain { ...Domain } }
+  ... on Chart     { tags { ...Tags } ownership { ...Owners } domain { ...Domain } }
+  ... on Dashboard { tags { ...Tags } ownership { ...Owners } domain { ...Domain } }
 }
 fragment Tags on GlobalTags {
   tags { tag { urn properties { name } } }
 }
+fragment Domain on DomainAssociation {
+  domain { urn properties { name description } }
+}
+fragment DataProduct on EntityRelationship {
+  entity {
+    urn
+    ... on DataProduct { properties { name description } }
+  }
+}
 fragment Owners on Ownership {
   owners {
     owner {
@@ -39,7 +55,7 @@ fragment Owners on Ownership {
     type
   }
 }
-";
+"#;
 
 /// Builds the request body for the [`RESOURCE_INFO_QUERY`], parameterized by the entity's URN.
 pub fn request(urn: &Urn) -> GraphQlRequest<'_> {
@@ -86,9 +102,12 @@ pub struct ResponseData {
 /// concrete entity type. [`Default`] yields the "no metadata" entity used when DataHub does not
 /// return an entity for a URN.
 #[derive(Debug, Default, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct Entity {
     tags: Option<GlobalTags>,
     ownership: Option<Ownership>,
+    domain: Option<DomainAssociation>,
+    data_products: Option<EntityRelationships>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -112,6 +131,47 @@ struct TagProperties {
     name: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct DomainAssociation {
+    domain: DomainNode,
+}
+
+#[derive(Debug, Deserialize)]
+struct DomainNode {
+    urn: Urn,
+    properties: Option<DomainProperties>,
+}
+
+#[derive(Debug, Deserialize)]
+struct DomainProperties {
+    name: String,
+    description: Option<String>,
+}
+
+/// The result of the `DataProductContains` relationship query. Each relationship's related entity is
+/// a data product the resource belongs to.
+#[derive(Debug, Deserialize)]
+struct EntityRelationships {
+    relationships: Vec<EntityRelationship>,
+}
+
+#[derive(Debug, Deserialize)]
+struct EntityRelationship {
+    entity: DataProductNode,
+}
+
+#[derive(Debug, Deserialize)]
+struct DataProductNode {
+    urn: Urn,
+    properties: Option<DataProductProperties>,
+}
+
+#[derive(Debug, Deserialize)]
+struct DataProductProperties {
+    name: String,
+    description: Option<String>,
+}
+
 #[derive(Debug, Deserialize)]
 struct Ownership {
     owners: Vec<Owner>,
@@ -189,6 +249,39 @@ impl Entity {
             })
             .collect();
 
+        let domain = self.domain.map(|association| {
+            let DomainNode { urn, properties } = association.domain;
+            // A domain without a properties aspect has no name; fall back to its URN.
+            let (name, description) = match properties {
+                Some(properties) => (properties.name, properties.description),
+                None => (urn.0.clone(), None),
+            };
+            Domain {
+                urn,
+                name,
+                description,
+            }
+        });
+
+        let data_products = self
+            .data_products
+            .into_iter()
+            .flat_map(|relationships| relationships.relationships)
+            .map(|relationship| {
+                let DataProductNode { urn, properties } = relationship.entity;
+                // A data product without a properties aspect has no name; fall back to its URN.
+                let (name, description) = match properties {
+                    Some(properties) => (properties.name, properties.description),
+                    None => (urn.0.clone(), None),
+                };
+                DataProduct {
+                    urn,
+                    name,
+                    description,
+                }
+            })
+            .collect();
+
         let mut owners: BTreeMap<Urn, Owners> = BTreeMap::new();
         for owner in self
             .ownership
@@ -210,7 +303,13 @@ impl Entity {
             }
         }
 
-        DataHubResourceInfoResponse { urn, tags, owners }
+        DataHubResourceInfoResponse {
+            urn,
+            tags,
+            domain,
+            data_products,
+            owners,
+        }
     }
 }
 

rust/resource-info-fetcher/src/backend/data_hub/mod.rs

@@ -91,6 +91,15 @@ pub struct DataHubResourceInfoResponse {
     urn: Urn,
     tags: Vec<Tag>,
 
+    /// The domain the resource belongs to, e.g. `urn:li:domain:marketing`. A resource is assigned to
+    /// at most one domain, so this is an [`Option`] rather than a list. [`None`] if the resource is
+    /// not assigned to any domain.
+    domain: Option<Domain>,
+
+    /// The data products the resource is part of, e.g. `urn:li:dataProduct:orders`. Modelled as a
+    /// list because DataHub allows an asset to belong to more than one data product.
+    data_products: Vec<DataProduct>,
+
     /// Owners grouped by their ownership type URN, e.g.
     /// `urn:li:ownershipType:__system__technical_owner`. DataHub has no fixed set of ownership
     /// types — users can define custom ones — so this is an open map, not an enum.
@@ -114,6 +123,22 @@ pub struct Tag {
     name: String,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Domain {
+    urn: Urn,
+    name: String,
+    description: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DataProduct {
+    urn: Urn,
+    name: String,
+    description: Option<String>,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct User {