diff --git a/ingestion/src/metadata/ingestion/source/database/unitycatalog/service_spec.py b/ingestion/src/metadata/ingestion/source/database/unitycatalog/service_spec.py index ca9f17b254e5..892f88ef7b74 100644 --- a/ingestion/src/metadata/ingestion/source/database/unitycatalog/service_spec.py +++ b/ingestion/src/metadata/ingestion/source/database/unitycatalog/service_spec.py @@ -11,6 +11,9 @@ from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import ( UnityCatalogProfilerInterface, ) +from metadata.profiler.interface.sqlalchemy.unity_catalog.sampler_interface import ( + UnityCatalogSamplerInterface, +) from metadata.utils.service_spec.default import DefaultDatabaseSpec ServiceSpec = DefaultDatabaseSpec( @@ -19,4 +22,5 @@ usage_source_class=UnitycatalogUsageSource, profiler_class=UnityCatalogProfilerInterface, test_suite_class=UnityCatalogTestSuiteInterface, + sampler_class=UnityCatalogSamplerInterface, ) diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py new file mode 100644 index 000000000000..12a4ae3eaacb --- /dev/null +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_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.ingestion.source.database.databricks.connection import ( + get_connection as databricks_get_connection, +) +from metadata.sampler.sqlalchemy.sampler import SQASampler + + +class UnityCatalogSamplerInterface(SQASampler): + def get_client(self): + """client is the session for SQA""" + self.connection = databricks_get_connection(self.service_connection_config) + self.client = super().get_client() + self.set_catalog(self.client) + + return self.client