Fix apps

open-metadata · Jan 31, 2024 · 49ca846 · 49ca846
1 parent 73226a9
commit 49ca846
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 22 deletions.
diff --git a/ingestion/src/metadata/applications/auto_tagger.py b/ingestion/src/metadata/applications/auto_tagger.py
@@ -21,6 +21,9 @@
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
+from metadata.generated.schema.metadataIngestion.application import (
+    OpenMetadataApplicationConfig,
+)
 from metadata.generated.schema.type.tagLabel import (
     LabelType,
     State,
@@ -60,16 +63,18 @@ class AutoTaggerApp(AppRunner):
           jwtToken: "..."
     """
 
-    def __init__(self, config: AutoTaggerAppConfig, metadata: OpenMetadata):
+    def __init__(self, config: OpenMetadataApplicationConfig, metadata: OpenMetadata):
         super().__init__(config, metadata)
 
-        if not isinstance(config, AutoTaggerAppConfig):
+        if not isinstance(self.app_config, AutoTaggerAppConfig):
             raise InvalidAppConfiguration(
                 f"AutoTagger Runner expects an AutoTaggerAppConfig, we got [{config}]"
             )
 
         self._ner_scanner = None
-        self.confidence_threshold = config.confidenceLevel or DEFAULT_CONFIDENCE
+        self.confidence_threshold = (
+            self.app_config.confidenceLevel or DEFAULT_CONFIDENCE
+        )
 
     @property
     def name(self) -> str:

diff --git a/ingestion/src/metadata/workflow/application.py b/ingestion/src/metadata/workflow/application.py
@@ -15,12 +15,6 @@
 from typing import List, Optional
 
 from metadata.config.common import WorkflowExecutionError
-from metadata.generated.schema.configuration.appsPrivateConfiguration import (
-    AppPrivateConfig,
-)
-from metadata.generated.schema.entity.applications.configuration.applicationConfig import (
-    AppConfig,
-)
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
@@ -54,12 +48,11 @@ class AppRunner(Step, ABC):
 
     def __init__(
         self,
-        config: AppConfig.__fields__["__root__"].type_,
-        private_config: AppPrivateConfig.__fields__["__root__"].type_,
+        config: OpenMetadataApplicationConfig,
         metadata: OpenMetadata,
     ):
-        self.config = config
-        self.private_config = private_config
+        self.app_config = config.appConfig
+        self.private_config = config.appPrivateConfig
         self.metadata = metadata
 
         super().__init__()
@@ -74,7 +67,8 @@ def run(self) -> None:
 
     @classmethod
     def create(cls, config_dict: dict, metadata: OpenMetadata) -> "Step":
-        return cls(config=config_dict, metadata=metadata)
+        config = OpenMetadataApplicationConfig.parse_obj(config_dict)
+        return cls(config=config, metadata=metadata)
 
 
 class ApplicationWorkflow(BaseWorkflow, ABC):
@@ -121,12 +115,7 @@ def post_init(self) -> None:
 
         try:
             self.runner = runner_class(
-                config=self.config.appConfig.__root__
-                if self.config.appConfig
-                else None,
-                private_config=self.config.appPrivateConfig.__root__
-                if self.config.appPrivateConfig
-                else None,
+                config=self.config,
                 metadata=self.metadata,
             )
         except Exception as exc:

diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py
@@ -76,10 +76,14 @@ def build_application_workflow_config(
         # We pass the generic class and let each app cast the actual object
         appConfig=AppConfig(
             __root__=application_pipeline_conf.appConfig.__root__,
-        ),
+        )
+        if application_pipeline_conf.appConfig
+        else None,
         appPrivateConfig=PrivateConfig(
             __root__=application_pipeline_conf.appPrivateConfig.__root__
-        ),
+        )
+        if application_pipeline_conf.appPrivateConfig
+        else None,
         workflowConfig=build_workflow_config_property(ingestion_pipeline),
         ingestionPipelineFQN=ingestion_pipeline.fullyQualifiedName.__root__,
     )