adding subtupes for openapi

Author: gabe-lyons

metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py

@@ -25,6 +25,7 @@ class DatasetSubTypes(StrEnum):
     NEO4J_NODE = "Neo4j Node"
     NEO4J_RELATIONSHIP = "Neo4j Relationship"
     SNOWFLAKE_STREAM = "Snowflake Stream"
+    API_ENDPOINT = "API Endpoint"
 
     # TODO: Create separate entity...
     NOTEBOOK = "Notebook"

metadata-ingestion/src/datahub/ingestion/source/identity/okta.py

@@ -8,9 +8,6 @@
 from typing import Dict, Iterable, List, Optional, Set, Union
 
 import nest_asyncio
-from okta.client import Client as OktaClient
-from okta.exceptions import OktaAPIException
-from okta.models import Group, GroupProfile, User, UserProfile, UserStatus
 from pydantic import validator
 from pydantic.fields import Field
 
@@ -50,6 +47,9 @@
     StatusClass,
 )
 from datahub.utilities.lossy_collections import LossyList
+from okta.client import Client as OktaClient
+from okta.exceptions import OktaAPIException
+from okta.models import Group, GroupProfile, User, UserProfile, UserStatus
 
 logger = logging.getLogger(__name__)
 nest_asyncio.apply()

metadata-ingestion/src/datahub/ingestion/source/openapi.py

@@ -20,6 +20,7 @@
 )
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.openapi_parser import (
     clean_url,
     compose_url_attr,
@@ -40,6 +41,7 @@
     GlobalTagsClass,
     InstitutionalMemoryClass,
     InstitutionalMemoryMetadataClass,
+    SubTypesClass,
     TagAssociationClass,
 )
 
@@ -262,6 +264,10 @@ def init_dataset(
         inst_memory = InstitutionalMemoryClass([link_metadata])
         dataset_snapshot.aspects.append(inst_memory)
 
+        # Add API endpoint subtype
+        sub_types = SubTypesClass(typeNames=[DatasetSubTypes.API_ENDPOINT])
+        dataset_snapshot.aspects.append(sub_types)
+
         return dataset_snapshot, dataset_name
 
     def build_wu(

metadata-ingestion/tests/integration/okta/test_okta.py

@@ -6,10 +6,10 @@
 import jsonpickle
 import pytest
 from freezegun import freeze_time
-from okta.models import Group, User
 
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.ingestion.source.identity.okta import OktaConfig
+from okta.models import Group, User
 from tests.test_helpers import mce_helpers
 from tests.test_helpers.state_helpers import (
     get_current_checkpoint_from_pipeline,