Fix #19807 - Implement the right SQA Sampler for UnityCatalog (#19811)

This commit is contained in:
Pere Miquel Brull 2025-02-16 20:58:27 +01:00 committed by GitHub
parent 0392e1a814
commit e4c34210ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 33 additions and 0 deletions

View File

@ -11,6 +11,9 @@ from metadata.ingestion.source.database.unitycatalog.usage import (
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 @@ ServiceSpec = DefaultDatabaseSpec(
usage_source_class=UnitycatalogUsageSource,
profiler_class=UnityCatalogProfilerInterface,
test_suite_class=UnityCatalogTestSuiteInterface,
sampler_class=UnityCatalogSamplerInterface,
)

View File

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