openmetadata-ingestion 1.2.3.1__py3-none-any.whl → 1.2.4.0__py3-none-any.whl

metadata/ingestion/api/parser.py

@@ -11,12 +11,12 @@
 """
 Helper to parse workflow configurations
 """
-from typing import Optional, Type, TypeVar, Union
+from typing import Type, TypeVar, Union
 
 from pydantic import BaseModel, ValidationError
 
-from metadata.generated.schema.entity.automations.testServiceConnection import (
-    TestServiceConnectionRequest,
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.dashboardService import (
     DashboardConnection,
@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseServiceType,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    IngestionPipeline,
+)
 from metadata.generated.schema.entity.services.messagingService import (
     MessagingConnection,
     MessagingServiceType,
@@ -271,13 +274,11 @@ def _parse_inner_connection(config_dict: dict, source_type: str) -> None:
     :param config_dict: JSON configuration
     :param source_type: source type name, e.g., Airflow.
     """
-    inner_source_type = config_dict["source"]["serviceConnection"]["config"][
-        "connection"
-    ]["type"]
+    inner_source_type = config_dict["type"]
     inner_service_type = get_service_type(inner_source_type)
     inner_connection_class = get_connection_class(inner_source_type, inner_service_type)
     _unsafe_parse_config(
-        config=config_dict["source"]["serviceConnection"]["config"]["connection"],
+        config=config_dict,
         cls=inner_connection_class,
         message=f"Error parsing the inner service connection for {source_type}",
     )
@@ -304,7 +305,12 @@ def parse_service_connection(config_dict: dict) -> None:
 
     if source_type in HAS_INNER_CONNECTION:
         # We will first parse the inner `connection` configuration
-        _parse_inner_connection(config_dict, source_type)
+        _parse_inner_connection(
+            config_dict["source"]["serviceConnection"]["config"]["connection"][
+                "config"
+            ]["connection"],
+            source_type,
+        )
 
     # Parse the service connection dictionary with the scoped class
     _unsafe_parse_config(
@@ -434,37 +440,84 @@ def parse_workflow_config_gracefully(
     raise ParsingConfigurationError("Uncaught error when parsing the workflow!")
 
 
-def parse_test_connection_request_gracefully(
+def parse_ingestion_pipeline_config_gracefully(
     config_dict: dict,
-) -> Optional[TestServiceConnectionRequest]:
+) -> IngestionPipeline:
     """
-    This function either correctly parses the pydantic class,
-    or throws a scoped error while fetching the required source
-    connection class
+    This function either correctly parses the pydantic class, or
+    throws a scoped error while fetching the required source connection
+    class.
 
-    :param config_dict: JSON workflow config
-    :return: TestServiceConnectionRequest or scoped error
+    :param config_dict: JSON ingestion pipeline config
+    :return:Ingestion Pipeline config or scoped error
     """
 
     try:
-        test_service_connection = TestServiceConnectionRequest.parse_obj(config_dict)
-        return test_service_connection
+        ingestion_pipeline = IngestionPipeline.parse_obj(config_dict)
+        return ingestion_pipeline
 
-    except ValidationError as err:
-        # Unsafe access to the keys. Allow a KeyError if the config is not well formatted
-        source_type = config_dict["connection"]["config"]["type"]
-        logger.warning(
-            f"Error parsing the Workflow Configuration for {source_type} ingestion: {err}"
+    except ValidationError:
+        source_config_type = config_dict["sourceConfig"]["config"].get("type")
+
+        if source_config_type is None:
+            raise InvalidWorkflowException("Missing type in the sourceConfig config")
+
+        source_config_class = get_source_config_class(source_config_type)
+
+        _unsafe_parse_config(
+            config=config_dict["sourceConfig"]["config"],
+            cls=source_config_class,
+            message="Error parsing the source config",
+        )
+
+    raise ParsingConfigurationError(
+        "Uncaught error when parsing the Ingestion Pipeline!"
+    )
+
+
+def parse_automation_workflow_gracefully(
+    config_dict: dict,
+) -> AutomationWorkflow:
+    """
+    This function either correctly parses the pydantic class, or
+    throws a scoped error while fetching the required source connection
+    class.
+
+    :param config_dict: JSON AutomationWorkflow config
+    :return: AutomationWorkflow config or scoped error
+    """
+
+    try:
+        automation_workflow = AutomationWorkflow.parse_obj(config_dict)
+        return automation_workflow
+
+    except ValidationError:
+        source_type = config_dict["request"]["connection"]["config"].get("type")
+
+        if source_type is None:
+            raise InvalidWorkflowException("Missing type in the connection config")
+
+        logger.debug(
+            f"Error parsing the Workflow Configuration for {source_type} ingestion"
         )
 
         service_type = get_service_type(source_type)
         connection_class = get_connection_class(source_type, service_type)
 
-        # Parse the dictionary with the scoped class
+        if source_type in HAS_INNER_CONNECTION:
+            # We will first parse the inner `connection` configuration
+            _parse_inner_connection(
+                config_dict["request"]["connection"]["config"]["connection"],
+                source_type,
+            )
+
+        # Parse the service connection dictionary with the scoped class
         _unsafe_parse_config(
-            config=config_dict["connection"]["config"],
+            config=config_dict["request"]["connection"]["config"],
             cls=connection_class,
-            message="Error parsing the connection config",
+            message="Error parsing the service connection",
         )
 
-    raise ParsingConfigurationError("Uncaught error when parsing the workflow!")
+    raise ParsingConfigurationError(
+        "Uncaught error when parsing the Ingestion Pipeline!"
+    )

metadata/ingestion/api/topology_runner.py

@@ -26,6 +26,7 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.ingestion.api.models import Either, Entity
+from metadata.ingestion.models.custom_properties import OMetaCustomProperties
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.patch_request import PatchRequest
 from metadata.ingestion.models.topology import (
@@ -387,6 +388,19 @@ class TopologyRunnerMixin(Generic[C]):
         # We'll keep the tag fqn in the context and use if required
         self.update_context(stage=stage, context=right)
 
+    @yield_and_update_context.register
+    def _(
+        self,
+        right: OMetaCustomProperties,
+        stage: NodeStage,
+        entity_request: Either[C],
+    ) -> Iterable[Either[Entity]]:
+        """Custom Property implementation for the context information"""
+        yield entity_request
+
+        # We'll keep the tag fqn in the context and use if required
+        self.update_context(stage=stage, context=right)
+
     @yield_and_update_context.register
     def _(
         self,
@@ -426,7 +440,7 @@ class TopologyRunnerMixin(Generic[C]):
         # Either use the received request or the acknowledged Entity
         entity = entity_request.right if entity_request else None
 
-        if not stage.nullable and entity is None:
+        if not stage.nullable and entity is None and entity_request.left is None:
             raise ValueError("Value unexpectedly None")
 
         if entity_request is not None:

metadata/ingestion/lineage/parser.py

@@ -117,26 +117,32 @@ class LineageParser:
         # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
         return self.retrieve_tables(self.parser.target_tables)
 
+    # pylint: disable=protected-access
     @cached_property
     def column_lineage(self) -> List[Tuple[Column, Column]]:
         """
         Get a list of tuples of column lineage
         """
-        if self.parser._dialect == SQLPARSE_DIALECT:  # pylint: disable=protected-access
-            return self.parser.get_column_lineage()
         column_lineage = []
-        for col_lineage in self.parser.get_column_lineage():
-            # In case of column level lineage it is possible that we get
-            # two or more columns as there might be some intermediate columns
-            # but the source columns will be the first value and
-            # the target column always will be the last columns
-            src_column = col_lineage[0]
-            tgt_column = col_lineage[-1]
-            src_col = Column(src_column.raw_name)
-            src_col._parent = src_column._parent  # pylint: disable=protected-access
-            tgt_col = Column(tgt_column.raw_name)
-            tgt_col._parent = tgt_column._parent  # pylint: disable=protected-access
-            column_lineage.append((src_col, tgt_col))
+        try:
+            if self.parser._dialect == SQLPARSE_DIALECT:
+                return self.parser.get_column_lineage()
+
+            for col_lineage in self.parser.get_column_lineage():
+                # In case of column level lineage it is possible that we get
+                # two or more columns as there might be some intermediate columns
+                # but the source columns will be the first value and
+                # the target column always will be the last columns
+                src_column = col_lineage[0]
+                tgt_column = col_lineage[-1]
+                src_col = Column(src_column.raw_name)
+                src_col._parent = src_column._parent  # pylint: disable=protected-access
+                tgt_col = Column(tgt_column.raw_name)
+                tgt_col._parent = tgt_column._parent  # pylint: disable=protected-access
+                column_lineage.append((src_col, tgt_col))
+        except Exception as err:
+            logger.warning(f"Failed to fetch column level lineage due to: {err}")
+            logger.debug(traceback.format_exc())
         return column_lineage
 
     @cached_property

metadata/ingestion/models/custom_properties.py

@@ -0,0 +1,63 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Custom models for custom properties
+"""
+from enum import Enum
+from typing import Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from metadata.generated.schema.api.data.createCustomProperty import (
+    CreateCustomPropertyRequest,
+)
+from metadata.generated.schema.type import basic, entityHistory
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class CustomPropertyDataTypes(Enum):
+    STRING = "string"
+    INTEGER = "integer"
+    MARKDOWN = "markdown"
+    DATE = "date"
+    DATETIME = "dateTime"
+    DURATION = "duration"
+    EMAIL = "email"
+    NUMBER = "number"
+    SQLQUERY = "sqlQuery"
+    TIME = "time"
+    TIMEINTERVAL = "timeInterval"
+    TIMESTAMP = "timestamp"
+
+
+class OMetaCustomProperties(BaseModel):
+    entity_type: Type[T]
+    custom_property_type: Optional[CustomPropertyDataTypes]
+    createCustomPropertyRequest: CreateCustomPropertyRequest
+
+
+class CustomPropertyType(BaseModel):
+    """
+    Pydantic Model for custom properties
+    """
+
+    id: basic.Uuid
+    name: basic.EntityName
+    displayName: Optional[str]
+    fullyQualifiedName: Optional[basic.FullyQualifiedEntityName]
+    description: Optional[basic.Markdown]
+    category: Optional[str]
+    nameSpace: Optional[str]
+    version: Optional[entityHistory.EntityVersion]
+    updatedAt: Optional[basic.Timestamp]
+    updatedBy: Optional[str]
+    href: Optional[basic.Href]

metadata/ingestion/ometa/mixins/custom_property_mixin.py

@@ -0,0 +1,80 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Mixin class containing Custom Property specific methods
+
+To be used by OpenMetadata class
+"""
+from typing import Dict
+
+from metadata.generated.schema.api.data.createCustomProperty import PropertyType
+from metadata.ingestion.models.custom_properties import (
+    CustomPropertyDataTypes,
+    CustomPropertyType,
+    OMetaCustomProperties,
+)
+from metadata.ingestion.ometa.client import REST
+from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaCustomPropertyMixin:
+    """
+    OpenMetadata API methods related to CustomProperty.
+
+    To be inherited by OpenMetadata
+    """
+
+    client: REST
+
+    def create_or_update_custom_property(
+        self, ometa_custom_property: OMetaCustomProperties
+    ) -> Dict:
+        """Create or update custom property. If custom property name matches an existing
+        one then it will be updated.
+
+        Args:
+            ometa_custom_property (OMetaCustomProperties): custom property to be create or updated
+        """
+        # Get the json schema id of the entity to be updated
+        entity_type = ENTITY_REFERENCE_TYPE_MAP.get(
+            ometa_custom_property.entity_type.__name__
+        )
+        entity_schema = self.client.get(
+            f"/metadata/types/name/{entity_type}?category=field"
+        )
+
+        # Get the data type of the custom property
+        if not ometa_custom_property.createCustomPropertyRequest.propertyType:
+            custom_property_type = self.get_custom_property_type(
+                data_type=ometa_custom_property.custom_property_type
+            )
+            property_type = PropertyType(id=custom_property_type.id, type="type")
+            ometa_custom_property.createCustomPropertyRequest.propertyType = (
+                property_type
+            )
+
+        resp = self.client.put(
+            f"/metadata/types/{entity_schema.get('id')}",
+            data=ometa_custom_property.createCustomPropertyRequest.json(),
+        )
+        return resp
+
+    def get_custom_property_type(
+        self, data_type: CustomPropertyDataTypes
+    ) -> CustomPropertyType:
+        """
+        Get all the supported datatypes for the custom properties
+        """
+        resp = self.client.get(f"/metadata/types/name/{data_type.value}?category=field")
+        return CustomPropertyType(**resp)

metadata/ingestion/ometa/mixins/ingestion_pipeline_mixin.py

@@ -20,6 +20,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
     IngestionPipeline,
     PipelineStatus,
 )
+from metadata.ingestion.api.parser import parse_ingestion_pipeline_config_gracefully
 from metadata.ingestion.ometa.client import REST
 from metadata.utils.logger import ometa_logger
 
@@ -79,7 +80,7 @@ class OMetaIngestionPipelineMixin:
             f"{self.get_suffix(IngestionPipeline)}/trigger/{ingestion_pipeline_id}"
         )
 
-        return IngestionPipeline.parse_obj(resp)
+        return parse_ingestion_pipeline_config_gracefully(resp)
 
     def get_pipeline_status_between_ts(
         self,
@@ -125,6 +126,6 @@ class OMetaIngestionPipelineMixin:
         )
 
         if hasattr(resp, "sourceConfig"):
-            return IngestionPipeline.parse_obj(resp)
+            return parse_ingestion_pipeline_config_gracefully(resp)
 
         return None

metadata/ingestion/ometa/mixins/patch_mixin_utils.py

@@ -21,10 +21,8 @@ from typing import Generic, List, Optional, Type, TypeVar, Union
 from pydantic import BaseModel
 
 from metadata.generated.schema.type import basic
-from metadata.utils.logger import ometa_logger
 
 T = TypeVar("T", bound=BaseModel)
-logger = ometa_logger()
 
 
 class PatchField(str, Enum):
@@ -140,9 +138,6 @@ class OMetaPatchMixinBase(Generic[T]):
         instance = self.get_by_id(entity=entity, entity_id=entity_id, fields=fields)
 
         if not instance:
-            logger.warning(
-                f"Cannot find an instance of '{entity.__class__.__name__}' with id [{str(entity_id)}]."
-            )
             return None
 
         return instance

metadata/ingestion/ometa/mixins/server_mixin.py

@@ -60,11 +60,13 @@ class OMetaServerMixin:
         Validate Server & Client versions. They should match.
         Otherwise, raise VersionMismatchException
         """
-        logger.debug("Validating client and server versions")
-
         server_version = self.get_server_version()
         client_version = get_client_version()
 
+        logger.info(
+            f"OpenMetadata client running with Server version [{server_version}] and Client version [{client_version}]"
+        )
+
         # Server version will be 0.13.2, vs 0.13.2.X from the client.
         # If the server version is contained in the client version, then we're good to go
         if server_version not in client_version:

metadata/ingestion/ometa/ometa_api.py

@@ -33,6 +33,9 @@ from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.models.encoders import show_secrets_encoder
 from metadata.ingestion.ometa.auth_provider import AuthenticationProvider
 from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
+from metadata.ingestion.ometa.mixins.custom_property_mixin import (
+    OMetaCustomPropertyMixin,
+)
 from metadata.ingestion.ometa.mixins.dashboard_mixin import OMetaDashboardMixin
 from metadata.ingestion.ometa.mixins.data_insight_mixin import DataInsightMixin
 from metadata.ingestion.ometa.mixins.es_mixin import ESMixin
@@ -110,6 +113,7 @@ class OpenMetadata(
     OMetaQueryMixin,
     OMetaRolePolicyMixin,
     OMetaSearchIndexMixin,
+    OMetaCustomPropertyMixin,
     Generic[T, C],
 ):
     """

metadata/ingestion/sink/metadata_rest.py

@@ -51,6 +51,7 @@ from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.generated.schema.type.schema import Topic
 from metadata.ingestion.api.models import Either, Entity, StackTraceError
 from metadata.ingestion.api.steps import Sink
+from metadata.ingestion.models.custom_properties import OMetaCustomProperties
 from metadata.ingestion.models.data_insight import OMetaDataInsightSample
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.life_cycle import OMetaLifeCycleData
@@ -172,6 +173,14 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         )
         return Either(right=entity)
 
+    @_run_dispatch.register
+    def write_custom_properties(self, record: OMetaCustomProperties) -> Either[Dict]:
+        """
+        Create or update the custom properties
+        """
+        custom_property = self.metadata.create_or_update_custom_property(record)
+        return Either(right=custom_property)
+
     @_run_dispatch.register
     def write_datamodel(self, datamodel_link: DataModelLink) -> Either[DataModel]:
         """

metadata/ingestion/source/dashboard/superset/api_source.py

@@ -123,19 +123,30 @@ class SupersetAPISource(SupersetSourceMixin):
     ) -> Iterable[Either[CreateChartRequest]]:
         """Method to fetch charts linked to dashboard"""
         for chart_id in self._get_charts_of_dashboard(dashboard_details):
-            chart_json = self.all_charts.get(chart_id)
-            if not chart_json:
-                logger.warning(f"chart details for id: {chart_id} not found, skipped")
-                continue
-            chart = CreateChartRequest(
-                name=chart_json.id,
-                displayName=chart_json.slice_name,
-                description=chart_json.description,
-                chartType=get_standard_chart_type(chart_json.viz_type),
-                sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{chart_json.url}",
-                service=self.context.dashboard_service,
-            )
-            yield Either(right=chart)
+            try:
+                chart_json = self.all_charts.get(chart_id)
+                if not chart_json:
+                    logger.warning(
+                        f"chart details for id: {chart_id} not found, skipped"
+                    )
+                    continue
+                chart = CreateChartRequest(
+                    name=chart_json.id,
+                    displayName=chart_json.slice_name,
+                    description=chart_json.description,
+                    chartType=get_standard_chart_type(chart_json.viz_type),
+                    sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{chart_json.url}",
+                    service=self.context.dashboard_service,
+                )
+                yield Either(right=chart)
+            except Exception as exc:  # pylint: disable=broad-except
+                yield Either(
+                    left=StackTraceError(
+                        name=chart_json.id,
+                        error=f"Error creating chart [{chart_json.id} - {chart_json.slice_name}]: {exc}",
+                        stack_trace=traceback.format_exc(),
+                    )
+                )
 
     def _get_datasource_fqn(
         self, datasource_id: str, db_service_entity: DatabaseService
@@ -166,7 +177,7 @@ class SupersetAPISource(SupersetSourceMixin):
                         service_name=db_service_entity.name.__root__,
                     )
                 return dataset_fqn
-        except KeyError as err:
+        except Exception as err:
             logger.debug(traceback.format_exc())
             logger.warning(
                 f"Failed to fetch Datasource with id [{datasource_id}]: {err}"