use common subtypes constants

datahub-project · Mar 5, 2025 · adb9679 · adb9679
1 parent db089a1
commit adb9679
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py b/metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py
@@ -60,8 +60,15 @@ class BIContainerSubTypes(StrEnum):
     MODE_COLLECTION = "Collection"
 
 
+class FlowContainerSubTypes(StrEnum):
+    MSSQL_JOB = "Job"
+    MSSQL_PROCEDURE_CONTAINER = "Procedures Container"
+
+
 class JobContainerSubTypes(StrEnum):
     NIFI_PROCESS_GROUP = "Process Group"
+    MSSQL_JOBSTEP = "Job Step"
+    MSSQL_STORED_PROCEDURE = "Stored Procedure"
 
 
 class BIAssetSubTypes(StrEnum):

diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/job_models.py
@@ -11,6 +11,10 @@
     DatabaseKey,
     SchemaKey,
 )
+from datahub.ingestion.source.common.subtypes import (
+    FlowContainerSubTypes,
+    JobContainerSubTypes,
+)
 from datahub.metadata.schema_classes import (
     ContainerClass,
     DataFlowInfoClass,
@@ -215,7 +219,11 @@ def as_datajob_info_aspect(self) -> DataJobInfoClass:
     @property
     def as_subtypes_aspect(self) -> SubTypesClass:
         assert isinstance(self.entity, (JobStep, StoredProcedure))
-        type = "Job Step" if isinstance(self.entity, JobStep) else "Stored Procedure"
+        type = (
+            JobContainerSubTypes.MSSQL_JOBSTEP
+            if isinstance(self.entity, JobStep)
+            else JobContainerSubTypes.MSSQL_STORED_PROCEDURE
+        )
         return SubTypesClass(
             typeNames=[type],
         )
@@ -288,7 +296,11 @@ def as_dataflow_info_aspect(self) -> DataFlowInfoClass:
     @property
     def as_subtypes_aspect(self) -> SubTypesClass:
         assert isinstance(self.entity, (MSSQLJob, MSSQLProceduresContainer))
-        type = "Job" if isinstance(self.entity, MSSQLJob) else "Procedures Container"
+        type = (
+            FlowContainerSubTypes.MSSQL_JOB
+            if isinstance(self.entity, MSSQLJob)
+            else FlowContainerSubTypes.MSSQL_PROCEDURE_CONTAINER
+        )
         return SubTypesClass(
             typeNames=[type],
         )