import uuid from typing import List from unittest import TestCase from metadata.generated.schema.api.data.createDashboardDataModel import ( CreateDashboardDataModelRequest, ) from metadata.generated.schema.api.data.createTable import CreateTableRequest from metadata.generated.schema.entity.data.dashboardDataModel import ( DashboardDataModel, DataModelType, ) from metadata.generated.schema.entity.data.table import ( Column, ColumnName, DataType, Table, TableConstraint, ) from metadata.generated.schema.type.basic import ( EntityExtension, EntityName, FullyQualifiedEntityName, Markdown, ) from metadata.generated.schema.type.entityReference import EntityReference from metadata.ingestion.models.custom_pydantic import BaseModel, CustomSecretStr class CustomPydanticValidationTest(TestCase): create_request = CreateTableRequest( name=EntityName("Sales::>Territory"), displayName="SalesTerritory", description=Markdown(root="Sales territory lookup table."), tableType="Regular", columns=[ Column( name=ColumnName(root="Sales::Last>Year"), displayName="SalesLastYear", dataType="NUMBER", arrayDataType=None, dataLength=None, precision=None, scale=None, dataTypeDisplay="NUMBER", description=Markdown(root="Sales total of previous year."), fullyQualifiedName=None, tags=None, constraint="NOT_NULL", ordinalPosition=7, jsonSchema=None, children=None, profile=None, customMetrics=None, ), Column( name=ColumnName(root="Bonus"), displayName="Bonus", dataType="NUMBER", arrayDataType=None, dataLength=None, precision=None, scale=None, dataTypeDisplay="NUMBER", description=Markdown(root="Bonus due if quota is met."), fullyQualifiedName=None, tags=None, constraint="NOT_NULL", ordinalPosition=4, jsonSchema=None, children=None, profile=None, customMetrics=None, ), Column( name=ColumnName(root="ModifiedDate"), displayName="ModifiedDate", dataType="DATETIME", arrayDataType=None, dataLength=None, precision=None, scale=None, dataTypeDisplay="DATETIME", description=Markdown(root="Date and time the record was last updated."), fullyQualifiedName=None, tags=None, constraint="NOT_NULL", ordinalPosition=9, jsonSchema=None, children=None, profile=None, customMetrics=None, ), ], dataModel=None, tableConstraints=[ TableConstraint(constraintType="PRIMARY_KEY", columns=["Sales::Last>Year"]) ], tablePartition=None, tableProfilerConfig=None, owners=None, databaseSchema=FullyQualifiedEntityName( root='New Gyro 360.New Gyro 360."AdventureWorks2017.HumanResources"' ), tags=None, schemaDefinition=None, retentionPeriod=None, extension=EntityExtension( root={ "DataQuality": '

Last evaluation: 07/24/2023
Interval: 30 days
Next run: 08/23/2023, 10:44:20
Measurement unit: percent [%]


MetricTargetLatest result

Completeness

90%
100%

Integrity

90%
100%

Timeliness

90%
25%

Uniqueness

90%
60%

Validity

90%
100%

Overall score of the table is: 77%

' } ), sourceUrl=None, domain=None, dataProducts=None, fileFormat=None, lifeCycle=None, sourceHash=None, ) create_request_dashboard_datamodel = CreateDashboardDataModelRequest( name=EntityName('test"dashboarddatamodel"'), displayName='test"dashboarddatamodel"', description=Markdown( root="test__reserved__quote__dashboarddatamodel__reserved__quote__" ), dataModelType=DataModelType.PowerBIDataModel, service=FullyQualifiedEntityName( root='New Gyro 360.New Gyro 360."AdventureWorks2017.HumanResources"' ), columns=[ Column( name="struct", dataType=DataType.STRUCT, arrayDataType="UNKNOWN", children=[ Column(name='test "struct_children"', dataType=DataType.BIGINT) ], ) ], ) def test_replace_separator(self): assert ( self.create_request.name.root == "Sales__reserved__colon____reserved__arrow__Territory" ) assert ( self.create_request.columns[0].name.root == "Sales__reserved__colon__Last__reserved__arrow__Year" ) assert ( self.create_request.tableConstraints[0].columns[0] == "Sales__reserved__colon__Last__reserved__arrow__Year" ) assert ( self.create_request_dashboard_datamodel.name.root == "test__reserved__quote__dashboarddatamodel__reserved__quote__" ) assert ( self.create_request_dashboard_datamodel.columns[0].children[0].name.root == "test __reserved__quote__struct_children__reserved__quote__" ) def test_revert_separator(self): fetch_response_revert_separator = Table( id=uuid.uuid4(), name="test__reserved__colon__table", databaseSchema=EntityReference(id=uuid.uuid4(), type="databaseSchema"), fullyQualifiedName="test-service-table.test-db.test-schema.test", columns=[Column(name="id", dataType=DataType.BIGINT)], ) fetch_response_revert_separator_2 = Table( id=uuid.uuid4(), name="test__reserved__colon__table__reserved__arrow__", databaseSchema=EntityReference(id=uuid.uuid4(), type="databaseSchema"), fullyQualifiedName="test-service-table.test-db.test-schema.test", columns=[Column(name="id", dataType=DataType.BIGINT)], ) fetch_response_revert_separator_3 = DashboardDataModel( id=uuid.uuid4(), name="test__reserved__quote__dashboarddatamodel__reserved__quote__", fullyQualifiedName="test-service-table.test-db.test-schema.test__reserved__quote__dashboarddatamodel__reserved__quote__", dataModelType=DataModelType.PowerBIDataModel, columns=[ Column( name="struct", dataType=DataType.STRUCT, children=[ Column(name='test "struct_children"', dataType=DataType.BIGINT) ], ) ], ) assert fetch_response_revert_separator_3.name.root == 'test"dashboarddatamodel"' assert ( fetch_response_revert_separator_3.columns[0].children[0].name.root == 'test "struct_children"' ) assert fetch_response_revert_separator.name.root == "test::table" assert fetch_response_revert_separator_2.name.root == "test::table>" class NestedModel(BaseModel): secret: CustomSecretStr value: int class RootModel(BaseModel): root_secret: CustomSecretStr nested: NestedModel items: List[NestedModel] data = { "root_secret": "root_password", "nested": {"secret": "nested_password", "value": 42}, "items": [ {"secret": "item1_password", "value": 1}, {"secret": "item2_password", "value": 2}, ], } model = RootModel(**data) masked_data = model.model_dump(mask_secrets=True) def test_model_dump_secrets(): """Test model_dump_masked with root, nested, and list structures.""" assert masked_data["root_secret"] == "**********" assert masked_data["nested"]["secret"] == "**********" assert masked_data["nested"]["value"] == 42 assert masked_data["items"][0]["secret"] == "**********" assert masked_data["items"][0]["value"] == 1 assert masked_data["items"][1]["secret"] == "**********" assert masked_data["items"][1]["value"] == 2 plain_data = model.model_dump(mask_secrets=False) assert plain_data["root_secret"] == "root_password" assert plain_data["nested"]["secret"] == "nested_password" assert plain_data["items"][0]["secret"] == "item1_password" default_dump = model.model_dump() assert default_dump["root_secret"] == "root_password" assert default_dump["nested"]["secret"] == "nested_password" assert default_dump["items"][0]["secret"] == "item1_password" def test_model_dump_json_secrets(): assert ( model.model_validate_json( model.model_dump_json() ).root_secret.get_secret_value() == "**********" ) assert ( model.model_validate_json( model.model_dump_json(mask_secrets=True) ).root_secret.get_secret_value() == "**********" ) assert ( model.model_validate_json( model.model_dump_json(mask_secrets=False) ).root_secret.get_secret_value() == "root_password" )