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

* Added custom properties methods in python sdk

* fixed tests
This commit is contained in:
Onkar Ravgan 2023-12-18 12:57:14 +05:30 committed by GitHub
parent d8984d267e
commit ef48b7eae7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 456 additions and 16 deletions

View File

@ -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,

View File

@ -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]

View File

@ -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)

View File

@ -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],
):
"""

View File

@ -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]:
"""

View File

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

View File

@ -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"])

View File

@ -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,20 +121,22 @@ public abstract class ServiceEntityRepository<
private void updateConnection() {
ServiceConnectionEntityInterface origConn = original.getConnection();
ServiceConnectionEntityInterface updatedConn = updated.getConnection();
String origJson = JsonUtils.pojoToJson(origConn);
String updatedJson = JsonUtils.pojoToJson(updatedConn);
S decryptedOrigConn = JsonUtils.readValue(origJson, serviceConnectionClass);
S decryptedUpdatedConn = JsonUtils.readValue(updatedJson, serviceConnectionClass);
SecretsManager secretsManager = SecretsManagerFactory.getSecretsManager();
decryptedOrigConn.setConfig(
secretsManager.decryptServiceConnectionConfig(
decryptedOrigConn.getConfig(), original.getServiceType().value(), serviceType));
decryptedUpdatedConn.setConfig(
secretsManager.decryptServiceConnectionConfig(
decryptedUpdatedConn.getConfig(), updated.getServiceType().value(), serviceType));
if (!objectMatch.test(decryptedOrigConn, decryptedUpdatedConn)) {
// we don't want save connection config details in our database
recordChange("connection", "old-encrypted-value", "new-encrypted-value", true);
if (!CommonUtil.nullOrEmpty(origConn) && !CommonUtil.nullOrEmpty(updatedConn)) {
String origJson = JsonUtils.pojoToJson(origConn);
String updatedJson = JsonUtils.pojoToJson(updatedConn);
S decryptedOrigConn = JsonUtils.readValue(origJson, serviceConnectionClass);
S decryptedUpdatedConn = JsonUtils.readValue(updatedJson, serviceConnectionClass);
SecretsManager secretsManager = SecretsManagerFactory.getSecretsManager();
decryptedOrigConn.setConfig(
secretsManager.decryptServiceConnectionConfig(
decryptedOrigConn.getConfig(), original.getServiceType().value(), serviceType));
decryptedUpdatedConn.setConfig(
secretsManager.decryptServiceConnectionConfig(
decryptedUpdatedConn.getConfig(), updated.getServiceType().value(), serviceType));
if (!objectMatch.test(decryptedOrigConn, decryptedUpdatedConn)) {
// we don't want save connection config details in our database
recordChange("connection", "old-encrypted-value", "new-encrypted-value", true);
}
}
}
}

View File

@ -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
}