import { Entity as EntityInterface, EntityType, SearchResult } from '../../types.generated'; import { FetchedEntity } from '../lineage/types'; import { Entity, EntityCapabilityType, IconStyleType, PreviewType } from './Entity'; import { GenericEntityProperties } from './shared/types'; import { dictToQueryStringParams, urlEncodeUrn } from './shared/utils'; function validatedGet(key: K, map: Map): V { if (map.has(key)) { return map.get(key) as V; } throw new Error(`Unrecognized key ${key} provided in map ${JSON.stringify(map)}`); } /** * Serves as a singleton registry for all DataHub entities to appear on the frontend. */ export default class EntityRegistry { entities: Array> = new Array>(); entityTypeToEntity: Map> = new Map>(); collectionNameToEntityType: Map = new Map(); pathNameToEntityType: Map = new Map(); register(entity: Entity) { this.entities.push(entity); this.entityTypeToEntity.set(entity.type, entity); this.collectionNameToEntityType.set(entity.getCollectionName(), entity.type); this.pathNameToEntityType.set(entity.getPathName(), entity.type); } getEntity(type: EntityType): Entity { return validatedGet(type, this.entityTypeToEntity); } getEntities(): Array> { return this.entities; } getSearchEntityTypes(): Array { return this.entities.filter((entity) => entity.isSearchEnabled()).map((entity) => entity.type); } getDefaultSearchEntityType(): EntityType { return this.entities[0].type; } getBrowseEntityTypes(): Array { return this.entities.filter((entity) => entity.isBrowseEnabled()).map((entity) => entity.type); } getLineageEntityTypes(): Array { return this.entities.filter((entity) => entity.isLineageEnabled()).map((entity) => entity.type); } getIcon(type: EntityType, fontSize: number, styleType: IconStyleType): JSX.Element { const entity = validatedGet(type, this.entityTypeToEntity); return entity.icon(fontSize, styleType); } getCollectionName(type: EntityType): string { const entity = validatedGet(type, this.entityTypeToEntity); return entity.getCollectionName(); } getEntityName(type: EntityType): string | undefined { const entity = validatedGet(type, this.entityTypeToEntity); return entity.getEntityName?.(); } getTypeFromCollectionName(name: string): EntityType { return validatedGet(name, this.collectionNameToEntityType); } getPathName(type: EntityType): string { const entity = validatedGet(type, this.entityTypeToEntity); return entity.getPathName(); } getEntityUrl(type: EntityType, urn: string, params?: Record): string { return `/${this.getPathName(type)}/${urlEncodeUrn(urn)}${params ? `?${dictToQueryStringParams(params)}` : ''}`; } getTypeFromPathName(pathName: string): EntityType { return validatedGet(pathName, this.pathNameToEntityType); } getTypeOrDefaultFromPathName(pathName: string, def?: EntityType): EntityType | undefined { try { return validatedGet(pathName, this.pathNameToEntityType); } catch (e) { return def; } } renderProfile(type: EntityType, urn: string): JSX.Element { const entity = validatedGet(type, this.entityTypeToEntity); return entity.renderProfile(urn); } renderPreview(entityType: EntityType, type: PreviewType, data: T): JSX.Element { const entity = validatedGet(entityType, this.entityTypeToEntity); return entity.renderPreview(type, data); } renderSearchResult(type: EntityType, searchResult: SearchResult): JSX.Element { const entity = validatedGet(type, this.entityTypeToEntity); return entity.renderSearch(searchResult); } renderBrowse(type: EntityType, data: T): JSX.Element { const entity = validatedGet(type, this.entityTypeToEntity); return entity.renderPreview(PreviewType.BROWSE, data); } getLineageVizConfig(type: EntityType, data: T): FetchedEntity | undefined { const entity = validatedGet(type, this.entityTypeToEntity); const genericEntityProperties = this.getGenericEntityProperties(type, data); return ( ({ ...entity.getLineageVizConfig?.(data), downstreamChildren: genericEntityProperties?.downstream?.relationships ?.filter((relationship) => relationship.entity) // eslint-disable-next-line @typescript-eslint/dot-notation ?.filter((relationship) => !relationship.entity?.['status']?.removed) ?.map((relationship) => ({ entity: relationship.entity as EntityInterface, type: (relationship.entity as EntityInterface).type, })), numDownstreamChildren: genericEntityProperties?.downstream?.total, upstreamChildren: genericEntityProperties?.upstream?.relationships ?.filter((relationship) => relationship.entity) // eslint-disable-next-line @typescript-eslint/dot-notation ?.filter((relationship) => !relationship.entity?.['status']?.removed) ?.map((relationship) => ({ entity: relationship.entity as EntityInterface, type: (relationship.entity as EntityInterface).type, })), numUpstreamChildren: genericEntityProperties?.upstream?.total, status: genericEntityProperties?.status, siblingPlatforms: genericEntityProperties?.siblingPlatforms, fineGrainedLineages: genericEntityProperties?.fineGrainedLineages, schemaMetadata: genericEntityProperties?.schemaMetadata, inputFields: genericEntityProperties?.inputFields, } as FetchedEntity) || undefined ); } getDisplayName(type: EntityType, data: T): string { const entity = validatedGet(type, this.entityTypeToEntity); return entity.displayName(data); } getGenericEntityProperties(type: EntityType, data: T): GenericEntityProperties | null { const entity = validatedGet(type, this.entityTypeToEntity); return entity.getGenericEntityProperties(data); } getSupportedEntityCapabilities(type: EntityType): Set { const entity = validatedGet(type, this.entityTypeToEntity); return entity.supportedCapabilities(); } getTypesWithSupportedCapabilities(capability: EntityCapabilityType): Set { return new Set( this.getEntities() .filter((entity) => entity.supportedCapabilities().has(capability)) .map((entity) => entity.type), ); } }