Keshav Mohta f1afe8f5f1
Fixes #23356: Databricks OAuth & Azure AD Auth (#23482)
* feat: databricks oauth and azure ad auth setup

* refactor: add auth type changes in databricks.md

* fix: test after oauth changes

* refactor: unity catalog connection to databricks connection code
2025-09-23 15:22:50 +05:30

49 lines
1.7 KiB
Python

# Copyright 2025 Collate
# Licensed under the Collate Community License, Version 1.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
# 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.generated.schema.entity.services.connections.database.databricksConnection import (
DatabricksConnection,
)
from metadata.sampler.sqlalchemy.databricks.sampler import DatabricksSamplerInterface
class UnityCatalogSamplerInterface(DatabricksSamplerInterface):
"""
Unity Catalog Sampler Interface
"""
def __init__(self, *args, **kwargs):
# Convert Unity Catalog connection to Databricks and move token to authType.
kwargs["service_connection_config"] = DatabricksConnection.model_validate(
{
**(
(
t := (
cfg := kwargs["service_connection_config"].model_dump(
mode="json"
)
).pop("token")
)
and cfg
),
"type": "Databricks",
"authType": {"token": t},
}
)
super().__init__(*args, **kwargs)