acryl-datahub 1.0.0.1rc6__py3-none-any.whl → 1.0.0.2__py3-none-any.whl

datahub/ingestion/source/superset.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 from dataclasses import dataclass, field
 from datetime import datetime
 from functools import lru_cache
@@ -100,6 +101,7 @@ from datahub.sql_parsing.sqlglot_lineage import (
 from datahub.utilities import config_clean
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.registries.domain_registry import DomainRegistry
+from datahub.utilities.threaded_iterator_executor import ThreadedIteratorExecutor
 
 logger = logging.getLogger(__name__)
 
@@ -210,6 +212,11 @@ class SupersetConfig(
         default=10, description="Timeout of single API call to superset."
     )
 
+    max_threads: int = Field(
+        default_factory=lambda: os.cpu_count() or 40,
+        description="Max parallelism for API calls. Defaults to cpuCount or 40",
+    )
+
     # TODO: Check and remove this if no longer needed.
     # Config database_alias is removed from sql sources.
     database_alias: Dict[str, str] = Field(
@@ -339,6 +346,7 @@ class SupersetSource(StatefulIngestionSourceBase):
 
             if response.status_code != 200:
                 logger.warning(f"Failed to get {entity_type} data: {response.text}")
+                continue
 
             payload = response.json()
             # Update total_items with the actual count from the response
@@ -501,33 +509,41 @@ class SupersetSource(StatefulIngestionSourceBase):
 
         return dashboard_snapshot
 
-    def emit_dashboard_mces(self) -> Iterable[MetadataWorkUnit]:
-        for dashboard_data in self.paginate_entity_api_results("dashboard/", PAGE_SIZE):
-            try:
-                dashboard_id = str(dashboard_data.get("id"))
-                dashboard_title = dashboard_data.get("dashboard_title", "")
-
-                if not self.config.dashboard_pattern.allowed(dashboard_title):
-                    self.report.report_dropped(
-                        f"Dashboard '{dashboard_title}' (id: {dashboard_id}) filtered by dashboard_pattern"
-                    )
-                    continue
-
-                dashboard_snapshot = self.construct_dashboard_from_api_data(
-                    dashboard_data
-                )
-            except Exception as e:
-                self.report.warning(
-                    f"Failed to construct dashboard snapshot. Dashboard name: {dashboard_data.get('dashboard_title')}. Error: \n{e}"
+    def _process_dashboard(self, dashboard_data: Any) -> Iterable[MetadataWorkUnit]:
+        dashboard_title = ""
+        try:
+            dashboard_id = str(dashboard_data.get("id"))
+            dashboard_title = dashboard_data.get("dashboard_title", "")
+            if not self.config.dashboard_pattern.allowed(dashboard_title):
+                self.report.report_dropped(
+                    f"Dashboard '{dashboard_title}' (id: {dashboard_id}) filtered by dashboard_pattern"
                 )
-                continue
-            # Emit the dashboard
-            mce = MetadataChangeEvent(proposedSnapshot=dashboard_snapshot)
-            yield MetadataWorkUnit(id=dashboard_snapshot.urn, mce=mce)
-            yield from self._get_domain_wu(
-                title=dashboard_title,
-                entity_urn=dashboard_snapshot.urn,
+                return
+            dashboard_snapshot = self.construct_dashboard_from_api_data(dashboard_data)
+        except Exception as e:
+            self.report.warning(
+                f"Failed to construct dashboard snapshot. Dashboard name: {dashboard_data.get('dashboard_title')}. Error: \n{e}"
+            )
+            return
+        mce = MetadataChangeEvent(proposedSnapshot=dashboard_snapshot)
+        yield MetadataWorkUnit(id=dashboard_snapshot.urn, mce=mce)
+        yield from self._get_domain_wu(
+            title=dashboard_title, entity_urn=dashboard_snapshot.urn
+        )
+
+    def emit_dashboard_mces(self) -> Iterable[MetadataWorkUnit]:
+        dashboard_data_list = [
+            (dashboard_data,)
+            for dashboard_data in self.paginate_entity_api_results(
+                "dashboard/", PAGE_SIZE
             )
+        ]
+
+        yield from ThreadedIteratorExecutor.process(
+            worker_func=self._process_dashboard,
+            args_list=dashboard_data_list,
+            max_workers=self.config.max_threads,
+        )
 
     def build_input_fields(
         self,
@@ -762,40 +778,46 @@ class SupersetSource(StatefulIngestionSourceBase):
             entity_urn=chart_urn,
         )
 
-    def emit_chart_mces(self) -> Iterable[MetadataWorkUnit]:
-        for chart_data in self.paginate_entity_api_results("chart/", PAGE_SIZE):
-            try:
-                chart_id = str(chart_data.get("id"))
-                chart_name = chart_data.get("slice_name", "")
-
-                if not self.config.chart_pattern.allowed(chart_name):
-                    self.report.report_dropped(
-                        f"Chart '{chart_name}' (id: {chart_id}) filtered by chart_pattern"
+    def _process_chart(self, chart_data: Any) -> Iterable[MetadataWorkUnit]:
+        chart_name = ""
+        try:
+            chart_id = str(chart_data.get("id"))
+            chart_name = chart_data.get("slice_name", "")
+            if not self.config.chart_pattern.allowed(chart_name):
+                self.report.report_dropped(
+                    f"Chart '{chart_name}' (id: {chart_id}) filtered by chart_pattern"
+                )
+                return
+            if self.config.dataset_pattern != AllowDenyPattern.allow_all():
+                datasource_id = chart_data.get("datasource_id")
+                if datasource_id:
+                    dataset_response = self.get_dataset_info(datasource_id)
+                    dataset_name = dataset_response.get("result", {}).get(
+                        "table_name", ""
                     )
-                    continue
-
-                # Emit a warning if charts use data from a dataset that will be filtered out
-                if self.config.dataset_pattern != AllowDenyPattern.allow_all():
-                    datasource_id = chart_data.get("datasource_id")
-                    if datasource_id:
-                        dataset_response = self.get_dataset_info(datasource_id)
-                        dataset_name = dataset_response.get("result", {}).get(
-                            "table_name", ""
+                    if dataset_name and not self.config.dataset_pattern.allowed(
+                        dataset_name
+                    ):
+                        self.report.warning(
+                            f"Chart '{chart_name}' (id: {chart_id}) uses dataset '{dataset_name}' which is filtered by dataset_pattern"
                         )
+            yield from self.construct_chart_from_chart_data(chart_data)
+        except Exception as e:
+            self.report.warning(
+                f"Failed to construct chart snapshot. Chart name: {chart_name}. Error: \n{e}"
+            )
+            return
 
-                        if dataset_name and not self.config.dataset_pattern.allowed(
-                            dataset_name
-                        ):
-                            self.report.warning(
-                                f"Chart '{chart_name}' (id: {chart_id}) uses dataset '{dataset_name}' which is filtered by dataset_pattern"
-                            )
-
-                yield from self.construct_chart_from_chart_data(chart_data)
-            except Exception as e:
-                self.report.warning(
-                    f"Failed to construct chart snapshot. Chart name: {chart_name}. Error: \n{e}"
-                )
-                continue
+    def emit_chart_mces(self) -> Iterable[MetadataWorkUnit]:
+        chart_data_list = [
+            (chart_data,)
+            for chart_data in self.paginate_entity_api_results("chart/", PAGE_SIZE)
+        ]
+        yield from ThreadedIteratorExecutor.process(
+            worker_func=self._process_chart,
+            args_list=chart_data_list,
+            max_workers=self.config.max_threads,
+        )
 
     def gen_schema_fields(self, column_data: List[Dict[str, str]]) -> List[SchemaField]:
         schema_fields: List[SchemaField] = []
@@ -1023,33 +1045,38 @@ class SupersetSource(StatefulIngestionSourceBase):
 
         return dataset_snapshot
 
-    def emit_dataset_mces(self) -> Iterable[MetadataWorkUnit]:
-        for dataset_data in self.paginate_entity_api_results("dataset/", PAGE_SIZE):
-            try:
-                dataset_name = dataset_data.get("table_name", "")
-
-                # Check if dataset should be filtered by dataset name
-                if not self.config.dataset_pattern.allowed(dataset_name):
-                    self.report.report_dropped(
-                        f"Dataset '{dataset_name}' filtered by dataset_pattern"
-                    )
-                    continue
-
-                dataset_snapshot = self.construct_dataset_from_dataset_data(
-                    dataset_data
-                )
-                mce = MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
-            except Exception as e:
-                self.report.warning(
-                    f"Failed to construct dataset snapshot. Dataset name: {dataset_data.get('table_name')}. Error: \n{e}"
+    def _process_dataset(self, dataset_data: Any) -> Iterable[MetadataWorkUnit]:
+        dataset_name = ""
+        try:
+            dataset_name = dataset_data.get("table_name", "")
+            if not self.config.dataset_pattern.allowed(dataset_name):
+                self.report.report_dropped(
+                    f"Dataset '{dataset_name}' filtered by dataset_pattern"
                 )
-                continue
-            # Emit the dataset
-            yield MetadataWorkUnit(id=dataset_snapshot.urn, mce=mce)
-            yield from self._get_domain_wu(
-                title=dataset_data.get("table_name", ""),
-                entity_urn=dataset_snapshot.urn,
+                return
+            dataset_snapshot = self.construct_dataset_from_dataset_data(dataset_data)
+            mce = MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
+        except Exception as e:
+            self.report.warning(
+                f"Failed to construct dataset snapshot. Dataset name: {dataset_data.get('table_name')}. Error: \n{e}"
             )
+            return
+        yield MetadataWorkUnit(id=dataset_snapshot.urn, mce=mce)
+        yield from self._get_domain_wu(
+            title=dataset_data.get("table_name", ""),
+            entity_urn=dataset_snapshot.urn,
+        )
+
+    def emit_dataset_mces(self) -> Iterable[MetadataWorkUnit]:
+        dataset_data_list = [
+            (dataset_data,)
+            for dataset_data in self.paginate_entity_api_results("dataset/", PAGE_SIZE)
+        ]
+        yield from ThreadedIteratorExecutor.process(
+            worker_func=self._process_dataset,
+            args_list=dataset_data_list,
+            max_workers=self.config.max_threads,
+        )
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         if self.config.ingest_dashboards:

datahub/ingestion/source/tableau/tableau.py

@@ -1623,7 +1623,7 @@ class TableauSiteSource:
                 # if multiple project has name C. Ideal solution is to use projectLuidWithin to avoid duplicate project,
                 # however Tableau supports projectLuidWithin in Tableau Cloud June 2022 / Server 2022.3 and later.
                 project_luid: Optional[str] = self._get_workbook_project_luid(workbook)
-                if project_luid not in self.tableau_project_registry.keys():
+                if project_luid not in self.tableau_project_registry:
                     wrk_name: Optional[str] = workbook.get(c.NAME)
                     wrk_id: Optional[str] = workbook.get(c.ID)
                     prj_name: Optional[str] = workbook.get(c.PROJECT_NAME)
@@ -2253,7 +2253,7 @@ class TableauSiteSource:
         # It is possible due to https://github.com/tableau/server-client-python/issues/1210
         if (
             ds.get(c.LUID)
-            and ds[c.LUID] not in self.datasource_project_map.keys()
+            and ds[c.LUID] not in self.datasource_project_map
             and self.report.get_all_datasources_query_failed
         ):
             logger.debug(
@@ -2265,7 +2265,7 @@ class TableauSiteSource:
 
         if (
             ds.get(c.LUID)
-            and ds[c.LUID] in self.datasource_project_map.keys()
+            and ds[c.LUID] in self.datasource_project_map
             and self.datasource_project_map[ds[c.LUID]] in self.tableau_project_registry
         ):
             return self.datasource_project_map[ds[c.LUID]]
@@ -3252,7 +3252,7 @@ class TableauSiteSource:
 
         parent_key = None
         project_luid: Optional[str] = self._get_workbook_project_luid(workbook)
-        if project_luid and project_luid in self.tableau_project_registry.keys():
+        if project_luid and project_luid in self.tableau_project_registry:
             parent_key = self.gen_project_key(project_luid)
         else:
             workbook_id: Optional[str] = workbook.get(c.ID)

datahub/ingestion/source/tableau/tableau_common.py

@@ -774,7 +774,7 @@ def get_overridden_info(
     if (
         lineage_overrides is not None
         and lineage_overrides.platform_override_map is not None
-        and original_platform in lineage_overrides.platform_override_map.keys()
+        and original_platform in lineage_overrides.platform_override_map
     ):
         platform = lineage_overrides.platform_override_map[original_platform]
 
@@ -782,7 +782,7 @@ def get_overridden_info(
         lineage_overrides is not None
         and lineage_overrides.database_override_map is not None
         and upstream_db is not None
-        and upstream_db in lineage_overrides.database_override_map.keys()
+        and upstream_db in lineage_overrides.database_override_map
     ):
         upstream_db = lineage_overrides.database_override_map[upstream_db]
 

datahub/ingestion/source/unity/source.py

@@ -1003,7 +1003,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             generate_usage_statistics=False,
             generate_operations=False,
         )
-        for dataset_name in self.view_definitions.keys():
+        for dataset_name in self.view_definitions:
             view_ref, view_definition = self.view_definitions[dataset_name]
             result = self._run_sql_parser(
                 view_ref,

datahub/ingestion/source/vertexai/vertexai.py

@@ -22,7 +22,11 @@ from google.oauth2 import service_account
 
 import datahub.emitter.mce_builder as builder
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.mcp_builder import ContainerKey, ProjectIdKey, gen_containers
+from datahub.emitter.mcp_builder import (
+    ExperimentKey,
+    ProjectIdKey,
+    gen_containers,
+)
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -96,10 +100,6 @@ class ModelMetadata:
     endpoints: Optional[List[Endpoint]] = None
 
 
-class ContainerKeyWithId(ContainerKey):
-    id: str
-
-
 @platform_name("Vertex AI", id="vertexai")
 @config_class(VertexAIConfig)
 @support_status(SupportStatus.TESTING)
@@ -173,7 +173,7 @@ class VertexAISource(Source):
     ) -> Iterable[MetadataWorkUnit]:
         yield from gen_containers(
             parent_container_key=self._get_project_container(),
-            container_key=ContainerKeyWithId(
+            container_key=ExperimentKey(
                 platform=self.platform,
                 id=self._make_vertexai_experiment_name(experiment.name),
             ),
@@ -309,7 +309,7 @@ class VertexAISource(Source):
     def _gen_experiment_run_mcps(
         self, experiment: Experiment, run: ExperimentRun
     ) -> Iterable[MetadataChangeProposalWrapper]:
-        experiment_key = ContainerKeyWithId(
+        experiment_key = ExperimentKey(
             platform=self.platform,
             id=self._make_vertexai_experiment_name(experiment.name),
         )

datahub/ingestion/transformer/add_dataset_dataproduct.py

@@ -54,7 +54,7 @@ class AddDatasetDataProduct(DatasetDataproductTransformer):
         data_products_container: Dict[str, DataProductPatchBuilder] = {}
         logger.debug("Generating dataproducts")
         is_container = self.config.is_container
-        for entity_urn in self.entity_map.keys():
+        for entity_urn in self.entity_map:
             data_product_urn = self.config.get_data_product_to_add(entity_urn)
             if data_product_urn:
                 if data_product_urn not in data_products:

datahub/ingestion/transformer/add_dataset_ownership.py

@@ -86,7 +86,7 @@ class AddDatasetOwnership(OwnershipTransformer):
         logger.debug("Generating Ownership for containers")
         ownership_container_mapping: Dict[str, List[OwnerClass]] = {}
         for entity_urn, data_ownerships in (
-            (urn, self.config.get_owners_to_add(urn)) for urn in self.entity_map.keys()
+            (urn, self.config.get_owners_to_add(urn)) for urn in self.entity_map
         ):
             if not data_ownerships:
                 continue

datahub/ingestion/transformer/dataset_domain.py

@@ -125,7 +125,7 @@ class AddDatasetDomain(DatasetDomainTransformer):
             return domain_mcps
 
         for entity_urn, domain_to_add in (
-            (urn, self.config.get_domains_to_add(urn)) for urn in self.entity_map.keys()
+            (urn, self.config.get_domains_to_add(urn)) for urn in self.entity_map
         ):
             if not domain_to_add or not domain_to_add.domains:
                 continue

datahub/lite/lite_util.py

@@ -99,7 +99,7 @@ def get_datahub_lite(config_dict: dict, read_only: bool = False) -> "DataHubLite
         lite_class = lite_registry.get(lite_type)
     except KeyError as e:
         raise Exception(
-            f"Failed to find a registered lite implementation for {lite_type}. Valid values are {[k for k in lite_registry.mapping.keys()]}"
+            f"Failed to find a registered lite implementation for {lite_type}. Valid values are {[k for k in lite_registry.mapping]}"
         ) from e
 
     lite_specific_config = lite_class.get_config_class().parse_obj(
@@ -127,7 +127,7 @@ def get_datahub_lite(config_dict: dict, read_only: bool = False) -> "DataHubLite
                 return lite
         else:
             raise Exception(
-                f"Failed to find a registered forwarding sink for type {lite_local_config.forward_to.type}. Valid values are {[k for k in sink_registry.mapping.keys()]}"
+                f"Failed to find a registered forwarding sink for type {lite_local_config.forward_to.type}. Valid values are {[k for k in sink_registry.mapping]}"
             )
     else:
         return lite

datahub/metadata/_schema_classes.py

@@ -15442,6 +15442,35 @@ class DataHubIngestionSourceKeyClass(_Aspect):
         self._inner_dict['id'] = value
     
     
+class DataHubOpenAPISchemaKeyClass(_Aspect):
+    """Key for a Query"""
+
+
+    ASPECT_NAME = 'dataHubOpenAPISchemaKey'
+    ASPECT_INFO = {'keyForEntity': 'dataHubOpenAPISchema', 'entityCategory': 'internal', 'entityAspects': ['systemMetadata'], 'entityDoc': 'Contains aspects which are used in OpenAPI requests/responses which are not otherwise present in the data model.'}
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.DataHubOpenAPISchemaKey")
+
+    def __init__(self,
+        id: str,
+    ):
+        super().__init__()
+        
+        self.id = id
+    
+    def _restore_defaults(self) -> None:
+        self.id = str()
+    
+    
+    @property
+    def id(self) -> str:
+        """A unique id for the DataHub OpenAPI schema."""
+        return self._inner_dict.get('id')  # type: ignore
+    
+    @id.setter
+    def id(self, value: str) -> None:
+        self._inner_dict['id'] = value
+    
+    
 class DataHubPersonaKeyClass(_Aspect):
     """Key for a persona type"""
 
@@ -20128,10 +20157,14 @@ class PlatformEventHeaderClass(DictWrapper):
         self._inner_dict['timestampMillis'] = value
     
     
-class SystemMetadataClass(DictWrapper):
+class SystemMetadataClass(_Aspect):
     """Metadata associated with each metadata change that is processed by the system"""
-    
+
+
+    ASPECT_NAME = 'systemMetadata'
+    ASPECT_INFO = {}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.mxe.SystemMetadata")
+
     def __init__(self,
         lastObserved: Optional[Union[int, None]]=None,
         runId: Optional[Union[str, None]]=None,
@@ -21738,6 +21771,9 @@ class QueryLanguageClass(object):
     SQL = "SQL"
     """A SQL Query"""
     
+    UNKNOWN = "UNKNOWN"
+    """Unknown query language"""
+    
     
     
 class QueryPropertiesClass(_Aspect):
@@ -26135,6 +26171,7 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.metadata.key.DataHubActionKey': DataHubActionKeyClass,
     'com.linkedin.pegasus2avro.metadata.key.DataHubConnectionKey': DataHubConnectionKeyClass,
     'com.linkedin.pegasus2avro.metadata.key.DataHubIngestionSourceKey': DataHubIngestionSourceKeyClass,
+    'com.linkedin.pegasus2avro.metadata.key.DataHubOpenAPISchemaKey': DataHubOpenAPISchemaKeyClass,
     'com.linkedin.pegasus2avro.metadata.key.DataHubPersonaKey': DataHubPersonaKeyClass,
     'com.linkedin.pegasus2avro.metadata.key.DataHubPolicyKey': DataHubPolicyKeyClass,
     'com.linkedin.pegasus2avro.metadata.key.DataHubRetentionKey': DataHubRetentionKeyClass,
@@ -26620,6 +26657,7 @@ __SCHEMA_TYPES = {
     'DataHubActionKey': DataHubActionKeyClass,
     'DataHubConnectionKey': DataHubConnectionKeyClass,
     'DataHubIngestionSourceKey': DataHubIngestionSourceKeyClass,
+    'DataHubOpenAPISchemaKey': DataHubOpenAPISchemaKeyClass,
     'DataHubPersonaKey': DataHubPersonaKeyClass,
     'DataHubPolicyKey': DataHubPolicyKeyClass,
     'DataHubRetentionKey': DataHubRetentionKeyClass,
@@ -26879,6 +26917,7 @@ ASPECT_CLASSES: List[Type[_Aspect]] = [
     ContainerClass,
     ContainerPropertiesClass,
     EditableContainerPropertiesClass,
+    SystemMetadataClass,
     DataHubSecretValueClass,
     DataHubUpgradeRequestClass,
     DataHubUpgradeResultClass,
@@ -26935,6 +26974,7 @@ ASPECT_CLASSES: List[Type[_Aspect]] = [
     MLModelKeyClass,
     NotebookKeyClass,
     RoleKeyClass,
+    DataHubOpenAPISchemaKeyClass,
     GlobalSettingsKeyClass,
     DatasetKeyClass,
     ChartKeyClass,
@@ -27102,6 +27142,7 @@ class AspectBag(TypedDict, total=False):
     container: ContainerClass
     containerProperties: ContainerPropertiesClass
     editableContainerProperties: EditableContainerPropertiesClass
+    systemMetadata: SystemMetadataClass
     dataHubSecretValue: DataHubSecretValueClass
     dataHubUpgradeRequest: DataHubUpgradeRequestClass
     dataHubUpgradeResult: DataHubUpgradeResultClass
@@ -27158,6 +27199,7 @@ class AspectBag(TypedDict, total=False):
     mlModelKey: MLModelKeyClass
     notebookKey: NotebookKeyClass
     roleKey: RoleKeyClass
+    dataHubOpenAPISchemaKey: DataHubOpenAPISchemaKeyClass
     globalSettingsKey: GlobalSettingsKeyClass
     datasetKey: DatasetKeyClass
     chartKey: ChartKeyClass
@@ -27292,6 +27334,7 @@ KEY_ASPECTS: Dict[str, Type[_Aspect]] = {
     'mlModel': MLModelKeyClass,
     'notebook': NotebookKeyClass,
     'role': RoleKeyClass,
+    'dataHubOpenAPISchema': DataHubOpenAPISchemaKeyClass,
     'globalSettings': GlobalSettingsKeyClass,
     'dataset': DatasetKeyClass,
     'chart': ChartKeyClass,
@@ -27352,6 +27395,7 @@ ENTITY_TYPE_NAMES: List[str] = [
     'mlModel',
     'notebook',
     'role',
+    'dataHubOpenAPISchema',
     'globalSettings',
     'dataset',
     'chart',
@@ -27411,6 +27455,7 @@ EntityTypeName = Literal[
     'mlModel',
     'notebook',
     'role',
+    'dataHubOpenAPISchema',
     'globalSettings',
     'dataset',
     'chart',

datahub/metadata/_urns/urn_defs.py

@@ -594,6 +594,62 @@ class RoleUrn(_SpecificUrn):
     def id(self) -> str:
         return self._entity_ids[0]
 
+if TYPE_CHECKING:
+    from datahub.metadata.schema_classes import DataHubOpenAPISchemaKeyClass
+
+class DataHubOpenAPISchemaUrn(_SpecificUrn):
+    ENTITY_TYPE: ClassVar[Literal["dataHubOpenAPISchema"]] = "dataHubOpenAPISchema"
+    _URN_PARTS: ClassVar[int] = 1
+
+    def __init__(self, id: Union["DataHubOpenAPISchemaUrn", str], *, _allow_coercion: bool = True) -> None:
+        if _allow_coercion:
+            # Field coercion logic (if any is required).
+            if isinstance(id, str):
+                if id.startswith('urn:li:'):
+                    try:
+                        id = DataHubOpenAPISchemaUrn.from_string(id)
+                    except InvalidUrnError:
+                        raise InvalidUrnError(f'Expecting a DataHubOpenAPISchemaUrn but got {id}')
+                else:
+                    id = UrnEncoder.encode_string(id)
+
+        # Validation logic.
+        if not id:
+            raise InvalidUrnError("DataHubOpenAPISchemaUrn id cannot be empty")
+        if isinstance(id, DataHubOpenAPISchemaUrn):
+            id = id.id
+        elif isinstance(id, Urn):
+            raise InvalidUrnError(f'Expecting a DataHubOpenAPISchemaUrn but got {id}')
+        if UrnEncoder.contains_reserved_char(id):
+            raise InvalidUrnError(f'DataHubOpenAPISchemaUrn id contains reserved characters')
+
+        super().__init__(self.ENTITY_TYPE, [id])
+
+    @classmethod
+    def _parse_ids(cls, entity_ids: List[str]) -> "DataHubOpenAPISchemaUrn":
+        if len(entity_ids) != cls._URN_PARTS:
+            raise InvalidUrnError(f"DataHubOpenAPISchemaUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
+        return cls(id=entity_ids[0], _allow_coercion=False)
+
+    @classmethod
+    def underlying_key_aspect_type(cls) -> Type["DataHubOpenAPISchemaKeyClass"]:
+        from datahub.metadata.schema_classes import DataHubOpenAPISchemaKeyClass
+
+        return DataHubOpenAPISchemaKeyClass
+
+    def to_key_aspect(self) -> "DataHubOpenAPISchemaKeyClass":
+        from datahub.metadata.schema_classes import DataHubOpenAPISchemaKeyClass
+
+        return DataHubOpenAPISchemaKeyClass(id=self.id)
+
+    @classmethod
+    def from_key_aspect(cls, key_aspect: "DataHubOpenAPISchemaKeyClass") -> "DataHubOpenAPISchemaUrn":
+        return cls(id=key_aspect.id)
+
+    @property
+    def id(self) -> str:
+        return self._entity_ids[0]
+
 if TYPE_CHECKING:
     from datahub.metadata.schema_classes import GlobalSettingsKeyClass
 

datahub/metadata/com/linkedin/pegasus2avro/metadata/key/__init__.py

@@ -19,6 +19,7 @@ from ......schema_classes import DataHubAccessTokenKeyClass
 from ......schema_classes import DataHubActionKeyClass
 from ......schema_classes import DataHubConnectionKeyClass
 from ......schema_classes import DataHubIngestionSourceKeyClass
+from ......schema_classes import DataHubOpenAPISchemaKeyClass
 from ......schema_classes import DataHubPersonaKeyClass
 from ......schema_classes import DataHubPolicyKeyClass
 from ......schema_classes import DataHubRetentionKeyClass
@@ -72,6 +73,7 @@ DataHubAccessTokenKey = DataHubAccessTokenKeyClass
 DataHubActionKey = DataHubActionKeyClass
 DataHubConnectionKey = DataHubConnectionKeyClass
 DataHubIngestionSourceKey = DataHubIngestionSourceKeyClass
+DataHubOpenAPISchemaKey = DataHubOpenAPISchemaKeyClass
 DataHubPersonaKey = DataHubPersonaKeyClass
 DataHubPolicyKey = DataHubPolicyKeyClass
 DataHubRetentionKey = DataHubRetentionKeyClass