From ebc0a551e5c882f9d56bd0fe78e9f445bceb537c Mon Sep 17 00:00:00 2001 From: Ayush Shah Date: Wed, 20 Dec 2023 21:18:05 +0530 Subject: [PATCH] Fixes 12947: Add Support For DQ and Profiler in Databricks Unity Catalog (#14424) --- .../databricks/test_suite_interface.py | 29 +++++++ .../snowflake/test_suite_interface.py | 29 +++++++ .../sqlalchemy/sqa_test_suite_interface.py | 12 +-- .../unity_catalog/test_suite_interface.py | 35 +++++++++ .../interface/test_suite_interface_factory.py | 40 ++++++++++ .../metadata/mixins/sqalchemy/sqa_mixin.py | 21 +++-- .../interface/profiler_interface_factory.py | 49 ++++++++---- .../databricks/profiler_interface.py | 26 +++++++ .../sqlalchemy/profiler_interface.py | 5 +- .../snowflake/profiler_interface.py | 7 ++ .../unity_catalog/profiler_interface.py | 33 ++++++++ .../unit/profiler/test_profiler_interface.py | 77 +++++++++++++++++++ .../tests/unit/test_suite/test_factories.py | 45 +++++++++++ .../database/unityCatalogConnection.json | 6 +- .../ui/src/locale/languages/en-us.json | 4 +- 15 files changed, 386 insertions(+), 32 deletions(-) create mode 100644 ingestion/src/metadata/data_quality/interface/sqlalchemy/databricks/test_suite_interface.py create mode 100644 ingestion/src/metadata/data_quality/interface/sqlalchemy/snowflake/test_suite_interface.py create mode 100644 ingestion/src/metadata/data_quality/interface/sqlalchemy/unity_catalog/test_suite_interface.py create mode 100644 ingestion/src/metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py create mode 100644 ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/profiler_interface.py diff --git a/ingestion/src/metadata/data_quality/interface/sqlalchemy/databricks/test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/sqlalchemy/databricks/test_suite_interface.py new file mode 100644 index 00000000000..68003c10ebf --- /dev/null +++ b/ingestion/src/metadata/data_quality/interface/sqlalchemy/databricks/test_suite_interface.py @@ -0,0 +1,29 @@ +# 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. + +""" +Interfaces with database for all database engine +supporting sqlalchemy abstraction layer +""" + + +from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import ( + SQATestSuiteInterface, +) + + +class DatabricksTestSuiteInterface(SQATestSuiteInterface): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def create_session(self): + super().create_session() + self.set_catalog(self.session) diff --git a/ingestion/src/metadata/data_quality/interface/sqlalchemy/snowflake/test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/sqlalchemy/snowflake/test_suite_interface.py new file mode 100644 index 00000000000..d5c3d877bc1 --- /dev/null +++ b/ingestion/src/metadata/data_quality/interface/sqlalchemy/snowflake/test_suite_interface.py @@ -0,0 +1,29 @@ +# 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. + +""" +Interfaces with database for all database engine +supporting sqlalchemy abstraction layer +""" + + +from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import ( + SQATestSuiteInterface, +) + + +class SnowflakeTestSuiteInterface(SQATestSuiteInterface): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def create_session(self): + super().create_session() + self.set_session_tag(self.session) diff --git a/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py index 0939d7c9e7c..a0cec41d773 100644 --- a/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py +++ b/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py @@ -58,12 +58,7 @@ class SQATestSuiteInterface(SQAInterfaceMixin, TestSuiteInterface): self.ometa_client = ometa_client self.table_entity = table_entity self.service_connection_config = service_connection_config - self.session = create_and_bind_session( - get_connection(self.service_connection_config) - ) - self.set_session_tag(self.session) - self.set_catalog(self.session) - + self.create_session() self._table = self._convert_table_to_orm_object(sqa_metadata) ( @@ -75,6 +70,11 @@ class SQATestSuiteInterface(SQAInterfaceMixin, TestSuiteInterface): self._sampler = self._create_sampler() self._runner = self._create_runner() + def create_session(self): + self.session = create_and_bind_session( + get_connection(self.service_connection_config) + ) + @property def sample(self) -> Union[DeclarativeMeta, AliasedClass]: """_summary_ diff --git a/ingestion/src/metadata/data_quality/interface/sqlalchemy/unity_catalog/test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/sqlalchemy/unity_catalog/test_suite_interface.py new file mode 100644 index 00000000000..afb4e9f886b --- /dev/null +++ b/ingestion/src/metadata/data_quality/interface/sqlalchemy/unity_catalog/test_suite_interface.py @@ -0,0 +1,35 @@ +# 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. + +""" +Interfaces with database for all database engine +supporting sqlalchemy abstraction layer +""" + + +from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import ( + SQATestSuiteInterface, +) +from metadata.ingestion.connections.session import create_and_bind_session +from metadata.ingestion.source.database.databricks.connection import ( + get_connection as databricks_get_connection, +) + + +class UnityCatalogTestSuiteInterface(SQATestSuiteInterface): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def create_session(self): + self.session = create_and_bind_session( + databricks_get_connection(self.service_connection_config) + ) + self.set_catalog(self.session) diff --git a/ingestion/src/metadata/data_quality/interface/test_suite_interface_factory.py b/ingestion/src/metadata/data_quality/interface/test_suite_interface_factory.py index ea35bc69eb6..da73b780019 100644 --- a/ingestion/src/metadata/data_quality/interface/test_suite_interface_factory.py +++ b/ingestion/src/metadata/data_quality/interface/test_suite_interface_factory.py @@ -17,14 +17,32 @@ from logging import Logger from metadata.data_quality.interface.pandas.pandas_test_suite_interface import ( PandasTestSuiteInterface, ) +from metadata.data_quality.interface.sqlalchemy.databricks.test_suite_interface import ( + DatabricksTestSuiteInterface, +) +from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import ( + SnowflakeTestSuiteInterface, +) from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import ( SQATestSuiteInterface, ) +from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import ( + UnityCatalogTestSuiteInterface, +) from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface from metadata.generated.schema.entity.data.table import Table +from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( + DatabricksConnection, +) from metadata.generated.schema.entity.services.connections.database.datalakeConnection import ( DatalakeConnection, ) +from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import ( + SnowflakeConnection, +) +from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import ( + UnityCatalogConnection, +) from metadata.generated.schema.entity.services.databaseService import DatabaseConnection from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.utils.logger import test_suite_logger @@ -51,6 +69,17 @@ class TestSuiteInterfaceFactory: """ self._interface_type[interface_type] = interface + def register_many(self, interface_dict): + """ + Registers multiple profiler interfaces at once. + + Args: + interface_dict: A dictionary mapping connection class names (strings) to their + corresponding profiler interface classes. + """ + for interface_type, interface_class in interface_dict.items(): + self.register(interface_type, interface_class) + def create( self, service_connection_config: DatabaseConnection, @@ -86,3 +115,14 @@ class TestSuiteInterfaceFactory: test_suite_interface_factory = TestSuiteInterfaceFactory() + + +test_suite_interface = { + DatabaseConnection.__name__: SQATestSuiteInterface, + DatalakeConnection.__name__: PandasTestSuiteInterface, + SnowflakeConnection.__name__: SnowflakeTestSuiteInterface, + UnityCatalogConnection.__name__: UnityCatalogTestSuiteInterface, + DatabricksConnection.__name__: DatabricksTestSuiteInterface, +} + +test_suite_interface_factory.register_many(test_suite_interface) diff --git a/ingestion/src/metadata/mixins/sqalchemy/sqa_mixin.py b/ingestion/src/metadata/mixins/sqalchemy/sqa_mixin.py index 0cab60c3ccb..fee729d7371 100644 --- a/ingestion/src/metadata/mixins/sqalchemy/sqa_mixin.py +++ b/ingestion/src/metadata/mixins/sqalchemy/sqa_mixin.py @@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.connections.database.databricksCo from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import ( SnowflakeType, ) +from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import ( + UnityCatalogConnection, +) from metadata.ingestion.source.connections import get_connection from metadata.ingestion.source.database.snowflake.queries import ( SNOWFLAKE_SESSION_TAG_QUERY, @@ -84,17 +87,21 @@ class SQAInterfaceMixin: ) def set_catalog(self, session) -> None: - """Set catalog for the session. Right now only databricks requires it + """Set catalog for the session. Right now only databricks and unity catalog requires it Args: session (Session): sqa session object """ - if isinstance(self.service_connection_config, DatabricksConnection): - bind = session.get_bind() - bind.execute( - "USE CATALOG %(catalog)s;", - {"catalog": self.service_connection_config.catalog}, - ).first() + if not isinstance( + self.service_connection_config, + (UnityCatalogConnection, DatabricksConnection), + ): + return + bind = session.get_bind() + bind.execute( + "USE CATALOG %(catalog)s;", + {"catalog": self.service_connection_config.catalog}, + ).first() def close(self): """close session""" diff --git a/ingestion/src/metadata/profiler/interface/profiler_interface_factory.py b/ingestion/src/metadata/profiler/interface/profiler_interface_factory.py index e969aaf3b02..10c30f810e9 100644 --- a/ingestion/src/metadata/profiler/interface/profiler_interface_factory.py +++ b/ingestion/src/metadata/profiler/interface/profiler_interface_factory.py @@ -18,6 +18,9 @@ from typing import cast from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import ( BigQueryConnection, ) +from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( + DatabricksConnection, +) from metadata.generated.schema.entity.services.connections.database.datalakeConnection import ( DatalakeConnection, ) @@ -30,6 +33,9 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon from metadata.generated.schema.entity.services.connections.database.trinoConnection import ( TrinoConnection, ) +from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import ( + UnityCatalogConnection, +) from metadata.generated.schema.entity.services.databaseService import DatabaseConnection from metadata.profiler.interface.pandas.profiler_interface import ( PandasProfilerInterface, @@ -38,6 +44,9 @@ from metadata.profiler.interface.profiler_interface import ProfilerInterface from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import ( BigQueryProfilerInterface, ) +from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import ( + DatabricksProfilerInterface, +) from metadata.profiler.interface.sqlalchemy.profiler_interface import ( SQAProfilerInterface, ) @@ -50,6 +59,9 @@ from metadata.profiler.interface.sqlalchemy.snowflake.profiler_interface import from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import ( TrinoProfilerInterface, ) +from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import ( + UnityCatalogProfilerInterface, +) class ProfilerInterfaceFactory: @@ -62,6 +74,17 @@ class ProfilerInterfaceFactory: """Register a new interface""" self._interface_type[interface_type] = interface_class + def register_many(self, interface_dict): + """ + Registers multiple profiler interfaces at once. + + Args: + interface_dict: A dictionary mapping connection class names (strings) to their + corresponding profiler interface classes. + """ + for interface_type, interface_class in interface_dict.items(): + self.register(interface_type, interface_class) + def create(self, interface_type: str, *args, **kwargs): """Create interface object based on interface type""" interface_class = self._interface_type.get(interface_type) @@ -72,17 +95,15 @@ class ProfilerInterfaceFactory: profiler_interface_factory = ProfilerInterfaceFactory() -profiler_interface_factory.register(DatabaseConnection.__name__, SQAProfilerInterface) -profiler_interface_factory.register( - BigQueryConnection.__name__, BigQueryProfilerInterface -) -profiler_interface_factory.register( - SingleStoreConnection.__name__, SingleStoreProfilerInterface -) -profiler_interface_factory.register( - DatalakeConnection.__name__, PandasProfilerInterface -) -profiler_interface_factory.register( - SnowflakeConnection.__name__, SnowflakeProfilerInterface -) -profiler_interface_factory.register(TrinoConnection.__name__, TrinoProfilerInterface) +profilers = { + DatabaseConnection.__name__: SQAProfilerInterface, + BigQueryConnection.__name__: BigQueryProfilerInterface, + SingleStoreConnection.__name__: SingleStoreProfilerInterface, + DatalakeConnection.__name__: PandasProfilerInterface, + SnowflakeConnection.__name__: SnowflakeProfilerInterface, + TrinoConnection.__name__: TrinoProfilerInterface, + UnityCatalogConnection.__name__: UnityCatalogProfilerInterface, + DatabricksConnection.__name__: DatabricksProfilerInterface, +} + +profiler_interface_factory.register_many(profilers) diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py new file mode 100644 index 00000000000..bbcbc85d7a9 --- /dev/null +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py @@ -0,0 +1,26 @@ +# 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. + +""" +Interfaces with database for all database engine +supporting sqlalchemy abstraction layer +""" + + +from metadata.profiler.interface.sqlalchemy.profiler_interface import ( + SQAProfilerInterface, +) + + +class DatabricksProfilerInterface(SQAProfilerInterface): + def __init__(self, service_connection_config, **kwargs): + super().__init__(service_connection_config=service_connection_config, **kwargs) + self.set_catalog(self.session) diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py index 968cbaf1579..1020e76d8c1 100644 --- a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py @@ -104,10 +104,11 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin): ) self._table = self._convert_table_to_orm_object(sqa_metadata) + self.create_session() + + def create_session(self): self.session_factory = self._session_factory() self.session = self.session_factory() - self.set_session_tag(self.session) - self.set_catalog(self.session) @property def table(self): diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/snowflake/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/snowflake/profiler_interface.py index ef3b3639402..0d4613fbc86 100644 --- a/ingestion/src/metadata/profiler/interface/sqlalchemy/snowflake/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/snowflake/profiler_interface.py @@ -29,6 +29,13 @@ class SnowflakeProfilerInterface(SQAProfilerInterface): sqlalchemy. """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def create_session(self): + super().create_session() + self.set_session_tag(self.session) + def _programming_error_static_metric(self, runner, column, exc, session, metrics): if exc.orig and exc.orig.errno in OVERFLOW_ERROR_CODES.get( session.bind.dialect.name diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/profiler_interface.py new file mode 100644 index 00000000000..148dd38e105 --- /dev/null +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/profiler_interface.py @@ -0,0 +1,33 @@ +# 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. + +""" +Interfaces with database for all database engine +supporting sqlalchemy abstraction layer +""" + + +from metadata.ingestion.source.database.databricks.connection import ( + get_connection as databricks_get_connection, +) +from metadata.profiler.interface.sqlalchemy.profiler_interface import ( + SQAProfilerInterface, +) + + +class UnityCatalogProfilerInterface(SQAProfilerInterface): + def __init__(self, service_connection_config, **kwargs): + super().__init__(service_connection_config=service_connection_config, **kwargs) + + def create_session(self): + self.connection = databricks_get_connection(self.service_connection_config) + super().create_session() + self.set_catalog(self.session) diff --git a/ingestion/tests/unit/profiler/test_profiler_interface.py b/ingestion/tests/unit/profiler/test_profiler_interface.py index be23b28694b..2dc878a8e37 100644 --- a/ingestion/tests/unit/profiler/test_profiler_interface.py +++ b/ingestion/tests/unit/profiler/test_profiler_interface.py @@ -30,6 +30,28 @@ from metadata.generated.schema.entity.services.connections.connectionBasicType i DataStorageConfig, SampleDataStorageConfig, ) +from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import ( + BigQueryConnection, +) +from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( + DatabricksConnection, +) +from metadata.generated.schema.entity.services.connections.database.datalakeConnection import ( + DatalakeConnection, +) +from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import ( + SingleStoreConnection, +) +from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import ( + SnowflakeConnection, +) +from metadata.generated.schema.entity.services.connections.database.trinoConnection import ( + TrinoConnection, +) +from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import ( + UnityCatalogConnection, +) +from metadata.generated.schema.entity.services.databaseService import DatabaseConnection from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import ( DatabaseServiceProfilerPipeline, ) @@ -40,7 +62,34 @@ from metadata.profiler.api.models import ( ProfileSampleConfig, TableConfig, ) +from metadata.profiler.interface.pandas.profiler_interface import ( + PandasProfilerInterface, +) from metadata.profiler.interface.profiler_interface import ProfilerInterface +from metadata.profiler.interface.profiler_interface_factory import ( + ProfilerInterfaceFactory, +) +from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import ( + BigQueryProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import ( + DatabricksProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.profiler_interface import ( + SQAProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface import ( + SingleStoreProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.snowflake.profiler_interface import ( + SnowflakeProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import ( + TrinoProfilerInterface, +) +from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import ( + UnityCatalogProfilerInterface, +) class ProfilerInterfaceTest(TestCase): @@ -305,3 +354,31 @@ class ProfilerInterfaceTest(TestCase): schema_config, table_fqn="demo" ), ) + + def test_register_many(self): + # Initialize factory + factory = ProfilerInterfaceFactory() + + # Define profiles dictionary + profiles = { + DatabaseConnection.__name__: SQAProfilerInterface, + BigQueryConnection.__name__: BigQueryProfilerInterface, + SingleStoreConnection.__name__: SingleStoreProfilerInterface, + DatalakeConnection.__name__: PandasProfilerInterface, + SnowflakeConnection.__name__: SnowflakeProfilerInterface, + TrinoConnection.__name__: TrinoProfilerInterface, + UnityCatalogConnection.__name__: UnityCatalogProfilerInterface, + DatabricksConnection.__name__: DatabricksProfilerInterface, + } + + # Register profiles + factory.register_many(profiles) + + # Assert all expected interfaces are registered + expected_interfaces = set(profiles.keys()) + actual_interfaces = set(factory._interface_type.keys()) + assert expected_interfaces == actual_interfaces + + # Assert profiler classes match registered interfaces + for interface_type, interface_class in profiles.items(): + assert factory._interface_type[interface_type] == interface_class diff --git a/ingestion/tests/unit/test_suite/test_factories.py b/ingestion/tests/unit/test_suite/test_factories.py index fbc492a7c56..cc4b1c0691a 100644 --- a/ingestion/tests/unit/test_suite/test_factories.py +++ b/ingestion/tests/unit/test_suite/test_factories.py @@ -18,12 +18,25 @@ from pytest import mark from metadata.data_quality.interface.pandas.pandas_test_suite_interface import ( PandasTestSuiteInterface, ) +from metadata.data_quality.interface.sqlalchemy.databricks.test_suite_interface import ( + DatabricksTestSuiteInterface, +) +from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import ( + SnowflakeTestSuiteInterface, +) from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import ( SQATestSuiteInterface, ) +from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import ( + UnityCatalogTestSuiteInterface, +) from metadata.data_quality.interface.test_suite_interface_factory import ( + TestSuiteInterfaceFactory, test_suite_interface_factory, ) +from metadata.generated.schema.entity.services.connections.database.databricksConnection import ( + DatabricksConnection, +) from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import ( S3Config, ) @@ -33,6 +46,12 @@ from metadata.generated.schema.entity.services.connections.database.datalakeConn from metadata.generated.schema.entity.services.connections.database.mysqlConnection import ( MysqlConnection, ) +from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import ( + SnowflakeConnection, +) +from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import ( + UnityCatalogConnection, +) from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials MYSQL_CONNECTION_CONFIG = MysqlConnection( @@ -73,3 +92,29 @@ def test_interface_factory( table_entity=None, # type: ignore ) assert interface.__class__ == expected_interface + + +def test_register_many(): + # Initialize factory + + factory = TestSuiteInterfaceFactory() + + test_suite_interfaces = { + "base": SQATestSuiteInterface, + DatalakeConnection.__name__: PandasTestSuiteInterface, + SnowflakeConnection.__name__: SnowflakeTestSuiteInterface, + UnityCatalogConnection.__name__: UnityCatalogTestSuiteInterface, + DatabricksConnection.__name__: DatabricksTestSuiteInterface, + } + + # Register profiles + factory.register_many(test_suite_interfaces) + + # Assert all expected interfaces are registered + expected_interfaces = set(test_suite_interfaces.keys()) + actual_interfaces = set(factory._interface_type.keys()) + assert expected_interfaces == actual_interfaces + + # Assert profiler classes match registered interfaces + for interface_type, interface_class in test_suite_interfaces.items(): + assert factory._interface_type[interface_type] == interface_class diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/unityCatalogConnection.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/unityCatalogConnection.json index 23113e43bcc..34008d498fc 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/unityCatalogConnection.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/unityCatalogConnection.json @@ -24,7 +24,7 @@ "title": "Service Type", "description": "Service Type", "$ref": "#/definitions/databricksType", - "default": "Databricks" + "default": "UnityCatalog" }, "scheme": { "title": "Connection Scheme", @@ -85,6 +85,10 @@ "title": "Supports Metadata Extraction", "$ref": "../connectionBasicType.json#/definitions/supportsMetadataExtraction" }, + "supportsProfiler": { + "title": "Supports Profiler", + "$ref": "../connectionBasicType.json#/definitions/supportsProfiler" + }, "supportsDatabase": { "title": "Supports Database", "$ref": "../connectionBasicType.json#/definitions/supportsDatabase" diff --git a/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json b/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json index d39f72523c6..7291c1aea00 100644 --- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json +++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json @@ -1177,8 +1177,8 @@ "message": { "access-block-time-message": "Access will be blocked for milliseconds after max failed login attapmts performed.", "access-to-collaborate": "Allow open access for anyone to join the team, view data, and collaborate.", - "action-has-been-done-but-deploy-successfully": "has been {{action}} and deployed successfully", - "action-has-been-done-but-failed-to-deploy": "has been {{action}}, but failed to deploy", + "action-has-been-done-but-deploy-successfully": " has been {{action}} and deployed successfully", + "action-has-been-done-but-failed-to-deploy": " has been {{action}}, but failed to deploy", "active-users": "Display the number of active users.", "add-data-asset-domain": "Start by adding a service or data asset to the {{domain}}.", "add-kpi-message": "Identify the Key Performance Indicators (KPI) that best reflect the health of your data assets. Review your data assets based on Description, Ownership, and Tier. Define your target metrics in absolute or percentage to track your progress. Finally, set a start and end date to achieve your data goals.",