MINOR: Added custom properties methods in python sdk (#14402)

* Added custom properties methods in python sdk

* fixed tests

ingestion/src/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,

ingestion/src/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]

ingestion/src/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)

ingestion/src/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
@@ -108,6 +111,7 @@ class OpenMetadata(
     OMetaQueryMixin,
     OMetaRolePolicyMixin,
     OMetaSearchIndexMixin,
+    OMetaCustomPropertyMixin,
     Generic[T, C],
 ):
     """

ingestion/src/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]:
         """

ingestion/src/metadata/utils/tag_utils.py

@@ -40,8 +40,8 @@ logger = ingestion_logger()
 def get_ometa_tag_and_classification(
     tags: List[str],
     classification_name: str,
-    tag_description: Optional[str],
+    tag_description: Optional[str] = None,
-    classification_description: Optional[str],
+    classification_description: Optional[str] = None,
     include_tags: bool = True,
     tag_fqn: Optional[FullyQualifiedEntityName] = None,
 ) -> Iterable[Either[OMetaTagAndClassification]]:

ingestion/tests/integration/ometa/test_ometa_custom_properties_api.py

@@ -0,0 +1,227 @@
+#  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.
+
+"""
+OpenMetadata high-level API Custom Properties Test
+"""
+from typing import Dict
+from unittest import TestCase
+
+from metadata.generated.schema.api.data.createCustomProperty import (
+    CreateCustomPropertyRequest,
+)
+from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
+from metadata.generated.schema.api.data.createDatabaseSchema import (
+    CreateDatabaseSchemaRequest,
+)
+from metadata.generated.schema.api.data.createTable import CreateTableRequest
+from metadata.generated.schema.api.services.createDatabaseService import (
+    CreateDatabaseServiceRequest,
+)
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.table import Column, DataType, Table
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseConnection,
+    DatabaseService,
+    DatabaseServiceType,
+)
+from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
+    OpenMetadataJWTClientConfig,
+)
+from metadata.generated.schema.type.basic import EntityExtension
+from metadata.ingestion.models.custom_properties import (
+    CustomPropertyDataTypes,
+    OMetaCustomProperties,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+
+class OMetaCustomAttributeTest(TestCase):
+    """
+    Run this integration test with the local API available
+    Install the ingestion package before running the tests
+    """
+
+    service_entity_id = None
+
+    server_config = OpenMetadataConnection(
+        hostPort="http://localhost:8585/api",
+        authProvider="openmetadata",
+        securityConfig=OpenMetadataJWTClientConfig(
+            jwtToken="eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
+        ),
+    )
+    metadata = OpenMetadata(server_config)
+
+    assert metadata.health_check()
+
+    service = CreateDatabaseServiceRequest(
+        name="test-service-custom-properties",
+        serviceType=DatabaseServiceType.Mysql,
+        connection=DatabaseConnection(
+            config=MysqlConnection(
+                username="username",
+                authType=BasicAuth(
+                    password="password",
+                ),
+                hostPort="http://localhost:1234",
+            )
+        ),
+    )
+    service_type = "databaseService"
+
+    def create_table(self, name: str, extensions: Dict) -> Table:
+        create = CreateTableRequest(
+            name=name,
+            databaseSchema=self.create_schema_entity.fullyQualifiedName,
+            columns=[Column(name="id", dataType=DataType.BIGINT)],
+            extension=EntityExtension(__root__=extensions),
+        )
+        return self.metadata.create_or_update(create)
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        """
+        Prepare ingredients
+        """
+
+        cls.service_entity = cls.metadata.create_or_update(data=cls.service)
+
+        create_db = CreateDatabaseRequest(
+            name="test-db",
+            service=cls.service_entity.fullyQualifiedName,
+        )
+
+        cls.create_db_entity = cls.metadata.create_or_update(data=create_db)
+
+        create_schema = CreateDatabaseSchemaRequest(
+            name="test-schema",
+            database=cls.create_db_entity.fullyQualifiedName,
+        )
+
+        cls.create_schema_entity = cls.metadata.create_or_update(data=create_schema)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        """
+        Clean up
+        """
+        service_id = str(
+            cls.metadata.get_by_name(
+                entity=DatabaseService, fqn="test-service-custom-properties"
+            ).id.__root__
+        )
+
+        cls.metadata.delete(
+            entity=DatabaseService,
+            entity_id=service_id,
+            recursive=True,
+            hard_delete=True,
+        )
+
+    def create_custom_property(self):
+        """
+        Test to create the custom property
+        """
+
+        # Create the table size property
+        ometa_custom_property_request = OMetaCustomProperties(
+            entity_type=Table,
+            custom_property_type=CustomPropertyDataTypes.STRING,
+            createCustomPropertyRequest=CreateCustomPropertyRequest(
+                name="TableSize", description="Size of the Table"
+            ),
+        )
+        self.metadata.create_or_update_custom_property(
+            ometa_custom_property=ometa_custom_property_request
+        )
+
+        # Create the DataQuality property for a table
+        ometa_custom_property_request = OMetaCustomProperties(
+            entity_type=Table,
+            custom_property_type=CustomPropertyDataTypes.MARKDOWN,
+            createCustomPropertyRequest=CreateCustomPropertyRequest(
+                name="DataQuality", description="Quality Details of a Table"
+            ),
+        )
+        self.metadata.create_or_update_custom_property(
+            ometa_custom_property=ometa_custom_property_request
+        )
+
+        # Create the SchemaCost property for database schema
+        ometa_custom_property_request = OMetaCustomProperties(
+            entity_type=DatabaseSchema,
+            custom_property_type=CustomPropertyDataTypes.INTEGER,
+            createCustomPropertyRequest=CreateCustomPropertyRequest(
+                name="SchemaAge", description="Age in years of a Schema"
+            ),
+        )
+        self.metadata.create_or_update_custom_property(
+            ometa_custom_property=ometa_custom_property_request
+        )
+
+    def test_add_custom_property_table(self):
+        """
+        Test to add the extension/custom property to the table
+        """
+
+        # create the custom properties
+        self.create_custom_property()
+
+        extensions = {
+            "DataQuality": '<div><p><b>Last evaluation:</b> 07/24/2023<br><b>Interval: </b>30 days <br><b>Next run:</b> 08/23/2023, 10:44:21<br><b>Measurement unit:</b> percent [%]</p><br><table><tbody><tr><th>Metric</th><th>Target</th><th>Latest result</th></tr><tr><td><p class="text-success">Completeness</p></td><td>90%</td><td><div class="bar fabric" style="width: 93%;"><strong>93%</strong></div></td></tr><tr><td><p class="text-success">Integrity</p></td><td>90%</td><td><div class="bar fabric" style="width: 100%;"><strong>100%</strong></div></td></tr><tr><td><p class="text-warning">Timeliness</p></td><td>90%</td><td><div class="bar fabric" style="width: 56%;"><strong>56%</strong></div></td></tr><tr><td><p class="text-success">Uniqueness</p></td><td>90%</td><td><div class="bar fabric" style="width: 100%;"><strong>100%</strong></div></td></tr><tr><td><p class="text-success">Validity</p></td><td>90%</td><td><div class="bar fabric" style="width: 100%;"><strong>100%</strong></div></td></tr></tbody></table><h3>Overall score of the table is: 89%</h3><hr style="border-width: 5px;"></div>',
+            "TableSize": "250 MB",
+        }
+
+        self.create_table(name="test_custom_properties", extensions=extensions)
+
+        res = self.metadata.get_by_name(
+            entity=Table,
+            fqn="test-service-custom-properties.test-db.test-schema.test_custom_properties",
+            fields=["*"],
+        )
+        self.assertEqual(
+            res.extension.__root__["DataQuality"], extensions["DataQuality"]
+        )
+        self.assertEqual(res.extension.__root__["TableSize"], extensions["TableSize"])
+
+    def test_add_custom_property_schema(self):
+        """
+        Test to add the extension/custom property to the schema
+        """
+
+        # create the custom properties
+        self.create_custom_property()
+
+        extensions = {"SchemaAge": 3}
+
+        create_schema = CreateDatabaseSchemaRequest(
+            name="test-schema-custom-property",
+            database=self.create_db_entity.fullyQualifiedName,
+            extension=EntityExtension(__root__=extensions),
+        )
+        self.metadata.create_or_update(data=create_schema)
+
+        res = self.metadata.get_by_name(
+            entity=DatabaseSchema,
+            fqn="test-service-custom-properties.test-db.test-schema-custom-property",
+            fields=["*"],
+        )
+        self.assertEqual(res.extension.__root__["SchemaAge"], extensions["SchemaAge"])

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ServiceEntityRepository.java

@@ -17,6 +17,7 @@ import static org.openmetadata.service.util.EntityUtil.objectMatch;
 import java.util.UUID;
 import lombok.Getter;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
+import org.openmetadata.common.utils.CommonUtil;
 import org.openmetadata.schema.ServiceConnectionEntityInterface;
 import org.openmetadata.schema.ServiceEntityInterface;
 import org.openmetadata.schema.entity.services.ServiceType;
@@ -120,6 +121,7 @@ public abstract class ServiceEntityRepository<
     private void updateConnection() {
       ServiceConnectionEntityInterface origConn = original.getConnection();
       ServiceConnectionEntityInterface updatedConn = updated.getConnection();
+      if (!CommonUtil.nullOrEmpty(origConn) && !CommonUtil.nullOrEmpty(updatedConn)) {
         String origJson = JsonUtils.pojoToJson(origConn);
         String updatedJson = JsonUtils.pojoToJson(updatedConn);
         S decryptedOrigConn = JsonUtils.readValue(origJson, serviceConnectionClass);
@@ -137,4 +139,5 @@ public abstract class ServiceEntityRepository<
         }
       }
     }
+  }
 }

openmetadata-spec/src/main/resources/json/schema/api/data/createCustomProperty.json

@@ -0,0 +1,40 @@
+{
+  "$id": "https://open-metadata.org/schema/api/data/createCustomProperty.json",
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "CreateCustomPropertyRequest",
+  "description": "Create Custom Property Model entity request",
+  "type": "object",
+  "definitions": {
+    "propertyType": {
+      "description": "Property Type",
+      "type": "object",
+      "properties": {
+        "type": {
+          "description": "Property type",
+          "type": "string",
+          "default": "type"
+        },
+        "id": {
+          "description": "Unique identifier of this instance.",
+          "$ref": "../../type/basic.json#/definitions/uuid"
+        }
+      }
+    }
+  },
+  "properties": {
+    "name": {
+      "description": "Name that identifies this Custom Property model.",
+      "$ref": "../../entity/data/container.json#/definitions/entityName"
+    },
+    "description": {
+      "description": "Description of the Container instance.",
+      "$ref": "../../type/basic.json#/definitions/markdown"
+    },
+    "propertyType": {
+      "description": "Property Type",
+      "$ref": "#/definitions/propertyType"
+    }
+  },
+  "required": ["name"],
+  "additionalProperties": false
+}