openmetadata-ingestion 1.9.13.0__py3-none-any.whl → 1.10.1.0__py3-none-any.whl

metadata/ingestion/source/database/unitycatalog/client.py

@@ -16,6 +16,20 @@ import traceback
 
 from requests import HTTPError
 
+from metadata.generated.schema.entity.services.connections.database.databricks.azureAdSetup import (
+    AzureAdSetup,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.databricksOAuth import (
+    DatabricksOauth,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
+    PersonalAccessToken,
+)
+from metadata.ingestion.source.database.databricks.auth import (
+    get_azure_ad_auth,
+    get_databricks_oauth_auth,
+    get_personal_access_token_auth,
+)
 from metadata.ingestion.source.database.databricks.client import (
     API_TIMEOUT,
     DatabricksClient,
@@ -27,7 +41,7 @@ from metadata.ingestion.source.database.unitycatalog.models import (
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
-TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage"
+TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage/get"
 COLUMN_LINEAGE_PATH = "/lineage-tracking/column-lineage/get"
 TABLES_PATH = "/unity-catalog/tables"
 
@@ -37,6 +51,26 @@ class UnityCatalogClient(DatabricksClient):
     UnityCatalogClient creates a Databricks connection based on DatabricksCredentials.
     """
 
+    def _get_auth_header(self) -> dict[str, str]:
+        """
+        Method to get auth header
+        """
+        auth_method = {
+            PersonalAccessToken: get_personal_access_token_auth,
+            DatabricksOauth: get_databricks_oauth_auth,
+            AzureAdSetup: get_azure_ad_auth,
+        }.get(type(self.config.authType))
+        if not auth_method:
+            raise ValueError(
+                f"Unsupported authentication type: {type(self.config.authType)}"
+            )
+
+        auth_args = auth_method(self.config)
+        if auth_args.get("access_token"):
+            return {"Authorization": f"Bearer {auth_args['access_token']}"}
+
+        return auth_args["credentials_provider"]()()
+
     def get_table_lineage(self, table_name: str) -> LineageTableStreams:
         """
         Method returns table lineage details

metadata/ingestion/source/database/unitycatalog/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from functools import partial
 from typing import Optional
 
@@ -22,6 +23,15 @@ from sqlalchemy.exc import DatabaseError
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.database.databricks.azureAdSetup import (
+    AzureAdSetup,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.databricksOAuth import (
+    DatabricksOauth,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
+    PersonalAccessToken,
+)
 from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
     UnityCatalogConnection,
 )
@@ -35,6 +45,7 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.databricks.auth import get_auth_config
 from metadata.ingestion.source.database.unitycatalog.models import DatabricksTable
 from metadata.ingestion.source.database.unitycatalog.queries import (
     UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
@@ -51,7 +62,7 @@ logger = ingestion_logger()
 
 
 def get_connection_url(connection: UnityCatalogConnection) -> str:
-    url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
+    url = f"{connection.scheme.value}://{connection.hostPort}"
     return url
 
 
@@ -59,10 +70,23 @@ def get_connection(connection: UnityCatalogConnection) -> WorkspaceClient:
     """
     Create connection
     """
+    client_params = {}
+    if isinstance(connection.authType, PersonalAccessToken):
+        client_params["token"] = connection.authType.token.get_secret_value()
+    elif isinstance(connection.authType, DatabricksOauth):
+        client_params["client_id"] = connection.authType.clientId
+        client_params[
+            "client_secret"
+        ] = connection.authType.clientSecret.get_secret_value()
+    elif isinstance(connection.authType, AzureAdSetup):
+        client_params["azure_client_id"] = connection.authType.azureClientId
+        client_params[
+            "azure_client_secret"
+        ] = connection.authType.azureClientSecret.get_secret_value()
+        client_params["azure_tenant_id"] = connection.authType.azureTenantId
 
     return WorkspaceClient(
-        host=get_host_from_host_port(connection.hostPort),
-        token=connection.token.get_secret_value(),
+        host=get_host_from_host_port(connection.hostPort), **client_params
     )
 
 
@@ -76,11 +100,19 @@ def get_sqlalchemy_connection(connection: UnityCatalogConnection) -> Engine:
             connection.connectionArguments = init_empty_connection_arguments()
         connection.connectionArguments.root["http_path"] = connection.httpPath
 
-    return create_generic_db_connection(
+    auth_args = get_auth_config(connection)
+
+    original_connection_arguments = connection.connectionArguments
+    connection.connectionArguments = deepcopy(original_connection_arguments)
+    connection.connectionArguments.root.update(auth_args)
+
+    engine = create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
+    connection.connectionArguments = original_connection_arguments
+    return engine
 
 
 def test_connection(

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -15,7 +15,6 @@ import traceback
 from typing import Iterable, Optional
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
-from metadata.generated.schema.entity.data.container import ContainerDataModel
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import Table
@@ -91,41 +90,6 @@ class UnitycatalogLineageSource(Source):
             )
         return cls(config, metadata)
 
-    def _get_data_model_column_fqn(
-        self, data_model_entity: ContainerDataModel, column: str
-    ) -> Optional[str]:
-        if not data_model_entity:
-            return None
-        for entity_column in data_model_entity.columns:
-            if entity_column.displayName.lower() == column.lower():
-                return entity_column.fullyQualifiedName.root
-        return None
-
-    def _get_container_column_lineage(
-        self, data_model_entity: ContainerDataModel, table_entity: Table
-    ) -> Optional[LineageDetails]:
-        try:
-            column_lineage = []
-            for column in table_entity.columns:
-                from_column = self._get_data_model_column_fqn(
-                    data_model_entity=data_model_entity, column=column.name.root
-                )
-                to_column = column.fullyQualifiedName.root
-                if from_column and to_column:
-                    column_lineage.append(
-                        ColumnLineage(fromColumns=[from_column], toColumn=to_column)
-                    )
-            if column_lineage:
-                return LineageDetails(
-                    columnsLineage=column_lineage,
-                    source=LineageSource.ExternalTableLineage,
-                )
-            return None
-        except Exception as exc:
-            logger.debug(f"Error computing container column lineage: {exc}")
-            logger.debug(traceback.format_exc())
-            return None
-
     def _get_lineage_details(
         self, from_table: Table, to_table: Table, databricks_table_fqn: str
     ) -> Optional[LineageDetails]:
@@ -160,87 +124,16 @@ class UnitycatalogLineageSource(Source):
             logger.debug(traceback.format_exc())
             return None
 
-    def _handle_external_location_lineage(
-        self, file_info, table: Table, is_upstream: bool
-    ) -> Iterable[Either[AddLineageRequest]]:
-        try:
-            if not file_info.storage_location:
-                logger.debug("No storage location found in fileInfo")
-                return
-
-            location_entity = self.metadata.es_search_container_by_path(
-                full_path=file_info.storage_location, fields="dataModel"
-            )
-
-            if location_entity and location_entity[0]:
-                lineage_details = None
-                if location_entity[0].dataModel:
-                    lineage_details = self._get_container_column_lineage(
-                        location_entity[0].dataModel, table
-                    )
-
-                if is_upstream:
-                    yield Either(
-                        left=None,
-                        right=AddLineageRequest(
-                            edge=EntitiesEdge(
-                                fromEntity=EntityReference(
-                                    id=location_entity[0].id,
-                                    type="container",
-                                ),
-                                toEntity=EntityReference(
-                                    id=table.id,
-                                    type="table",
-                                ),
-                                lineageDetails=lineage_details,
-                            )
-                        ),
-                    )
-                else:
-                    yield Either(
-                        left=None,
-                        right=AddLineageRequest(
-                            edge=EntitiesEdge(
-                                fromEntity=EntityReference(
-                                    id=table.id,
-                                    type="table",
-                                ),
-                                toEntity=EntityReference(
-                                    id=location_entity[0].id,
-                                    type="container",
-                                ),
-                                lineageDetails=lineage_details,
-                            )
-                        ),
-                    )
-            else:
-                logger.debug(
-                    f"Unable to find container for external location: {file_info.storage_location}"
-                )
-        except Exception as exc:
-            logger.debug(
-                f"Error while processing external location lineage for {file_info.storage_location}: {exc}"
-            )
-            logger.debug(traceback.format_exc())
-
     def _handle_upstream_table(
         self,
         table_streams: LineageTableStreams,
         table: Table,
         databricks_table_fqn: str,
     ) -> Iterable[Either[AddLineageRequest]]:
-        for upstream_entity in table_streams.upstreams:
+        for upstream_table in table_streams.upstream_tables:
             try:
-                if upstream_entity.fileInfo:
-                    yield from self._handle_external_location_lineage(
-                        upstream_entity.fileInfo, table, is_upstream=True
-                    )
-                    continue
-
-                if not upstream_entity.tableInfo or not upstream_entity.tableInfo.name:
+                if not upstream_table.name:
                     continue
-
-                upstream_table = upstream_entity.tableInfo
                 from_entity_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Table,
@@ -279,69 +172,9 @@ class UnitycatalogLineageSource(Source):
                     )
             except Exception:
                 logger.debug(
-                    "Error while processing upstream lineage for "
-                    f"{databricks_table_fqn}"
-                )
-                logger.debug(traceback.format_exc())
-
-    def _handle_downstream_table(
-        self,
-        table_streams: LineageTableStreams,
-        table: Table,
-        databricks_table_fqn: str,
-    ) -> Iterable[Either[AddLineageRequest]]:
-        for downstream_entity in table_streams.downstreams:
-            try:
-                if downstream_entity.fileInfo:
-                    yield from self._handle_external_location_lineage(
-                        downstream_entity.fileInfo, table, is_upstream=False
-                    )
-                    continue
-
-                if (
-                    not downstream_entity.tableInfo
-                    or not downstream_entity.tableInfo.name
-                ):
-                    continue
-
-                downstream_table = downstream_entity.tableInfo
-                to_entity_fqn = fqn.build(
-                    metadata=self.metadata,
-                    entity_type=Table,
-                    database_name=downstream_table.catalog_name,
-                    schema_name=downstream_table.schema_name,
-                    table_name=downstream_table.name,
-                    service_name=self.config.serviceName,
-                )
-
-                to_entity = self.metadata.get_by_name(entity=Table, fqn=to_entity_fqn)
-                if to_entity:
-                    downstream_table_fqn = f"{downstream_table.catalog_name}.{downstream_table.schema_name}.{downstream_table.name}"
-                    lineage_details = self._get_lineage_details(
-                        from_table=table,
-                        to_table=to_entity,
-                        databricks_table_fqn=downstream_table_fqn,
-                    )
-                    yield Either(
-                        left=None,
-                        right=AddLineageRequest(
-                            edge=EntitiesEdge(
-                                fromEntity=EntityReference(id=table.id, type="table"),
-                                toEntity=EntityReference(id=to_entity.id, type="table"),
-                                lineageDetails=lineage_details,
-                            )
-                        ),
-                    )
-                else:
-                    logger.debug(
-                        f"Unable to find downstream entity for "
-                        f"{databricks_table_fqn} -> "
-                        f"{downstream_table.catalog_name}.{downstream_table.schema_name}.{downstream_table.name}"
-                    )
-            except Exception:
-                logger.debug(
-                    "Error while processing downstream lineage for "
-                    f"{databricks_table_fqn}"
+                    "Error while processing lineage for "
+                    f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
+                    f" -> {databricks_table_fqn}"
                 )
                 logger.debug(traceback.format_exc())
 
@@ -392,17 +225,10 @@ class UnitycatalogLineageSource(Source):
                     table_streams: LineageTableStreams = self.client.get_table_lineage(
                         databricks_table_fqn
                     )
-
-                    # Process upstream lineage
                     yield from self._handle_upstream_table(
                         table_streams, table, databricks_table_fqn
                     )
 
-                    # Process downstream lineage
-                    yield from self._handle_downstream_table(
-                        table_streams, table, databricks_table_fqn
-                    )
-
     def test_connection(self) -> None:
         test_connection_common(
             self.metadata, self.connection_obj, self.service_connection

metadata/ingestion/source/database/unitycatalog/models.py

@@ -22,8 +22,6 @@ class DatabricksTable(BaseModel):
     name: Optional[str] = None
     catalog_name: Optional[str] = None
     schema_name: Optional[str] = None
-    table_type: Optional[str] = None
-    lineage_timestamp: Optional[str] = None
 
 
 class DatabricksColumn(BaseModel):
@@ -33,23 +31,9 @@ class DatabricksColumn(BaseModel):
     table_name: Optional[str] = None
 
 
-class FileInfo(BaseModel):
-    path: Optional[str] = None
-    has_permission: Optional[bool] = None
-    securable_name: Optional[str] = None
-    storage_location: Optional[str] = None
-    securable_type: Optional[str] = None
-    lineage_timestamp: Optional[str] = None
-
-
-class LineageEntity(BaseModel):
-    tableInfo: Optional[DatabricksTable] = None
-    fileInfo: Optional[FileInfo] = None
-
-
 class LineageTableStreams(BaseModel):
-    upstreams: Optional[List[LineageEntity]] = []
-    downstreams: Optional[List[LineageEntity]] = []
+    upstream_tables: Optional[List[DatabricksTable]] = []
+    downstream_tables: Optional[List[DatabricksTable]] = []
 
 
 class LineageColumnStreams(BaseModel):

metadata/ingestion/source/messaging/kafka/connection.py

@@ -16,8 +16,9 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from confluent_kafka import DeserializingConsumer
 from confluent_kafka.admin import AdminClient, KafkaException
-from confluent_kafka.avro import AvroConsumer
+from confluent_kafka.schema_registry.avro import AvroDeserializer
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
 
 from metadata.generated.schema.entity.automations.workflow import (
@@ -71,6 +72,7 @@ def get_connection(
     """
     consumer_config = deepcopy(connection.consumerConfig) or {}
     schema_registry_config = deepcopy(connection.schemaRegistryConfig) or {}
+
     if connection.saslUsername or connection.saslPassword or connection.saslMechanism:
         if connection.saslUsername:
             consumer_config["sasl.username"] = connection.saslUsername
@@ -98,19 +100,26 @@ def get_connection(
 
     schema_registry_client = None
     consumer_client = None
+
     if connection.schemaRegistryURL:
         schema_registry_config["url"] = str(connection.schemaRegistryURL)
         schema_registry_client = SchemaRegistryClient(schema_registry_config)
+
         consumer_config["bootstrap.servers"] = connection.bootstrapServers
         if "group.id" not in consumer_config:
             consumer_config["group.id"] = "openmetadata-consumer"
         if "auto.offset.reset" not in consumer_config:
             consumer_config["auto.offset.reset"] = "largest"
         consumer_config["enable.auto.commit"] = False
-        logger.debug(f"Using Kafka consumer config: {consumer_config}")
-        consumer_client = AvroConsumer(
-            consumer_config, schema_registry=schema_registry_client
+
+        avro_deserializer = AvroDeserializer(
+            schema_registry_client=schema_registry_client
         )
+        consumer_config["value.deserializer"] = avro_deserializer
+
+        logger.debug(f"Using Kafka consumer config: {consumer_config}")
+
+        consumer_client = DeserializingConsumer(consumer_config)
 
     return KafkaClient(
         admin_client=admin_client,

metadata/ingestion/source/metadata/amundsen/client.py

@@ -59,16 +59,16 @@ class Neo4jHelper:
         """
         Create a Neo4j connection to Database
         """
-        trust = (
-            neo4j.TRUST_SYSTEM_CA_SIGNED_CERTIFICATES
-            if self.conf.neo4j_validate_ssl
-            else neo4j.TRUST_ALL_CERTIFICATES
-        )
+        if self.conf.neo4j_validate_ssl:
+            trusted_certificates = neo4j.TrustSystemCAs()
+        else:
+            trusted_certificates = neo4j.TrustAll()
+
         return GraphDatabase.driver(
             self.graph_url,
             auth=(self.conf.username, self.conf.password),
             encrypted=self.conf.neo4j_encrypted,
-            trust=trust,
+            trusted_certificates=trusted_certificates,
         )
 
     def _execute_query(self, transaction: Any, query: str) -> Any:

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -12,12 +12,13 @@
 Base class for ingesting mlmodel services
 """
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional, Set
+from typing import Any, Iterable, List, Optional, Set, Tuple
 
 from pydantic import Field
 from typing_extensions import Annotated
 
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.mlmodel import (
     MlFeature,
     MlHyperParameter,
@@ -39,6 +40,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
@@ -92,6 +94,12 @@ class MlModelServiceTopology(ServiceTopology):
                 consumer=["mlmodel_service"],
                 use_cache=True,
             ),
+            NodeStage(
+                type_=AddLineageRequest,
+                processor="yield_mlmodel_lineage",
+                consumer=["mlmodel_service"],
+                nullable=True,
+            ),
         ],
     )
 
@@ -208,3 +216,69 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def prepare(self):
         """By default, nothing to prepare"""
+
+    def get_db_service_prefixes(self) -> List[str]:
+        """
+        Get the list of db service prefixes
+        """
+        return (
+            self.source_config.lineageInformation.dbServicePrefixes or []
+            if hasattr(self.source_config, "lineageInformation")
+            and self.source_config.lineageInformation
+            else []
+        )
+
+    def parse_db_service_prefix(
+        self, db_service_prefix: Optional[str]
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+        """
+        Parse the db service prefix
+        Returns:
+            Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]: service, database, schema, table
+        """
+        prefix_parts = (db_service_prefix or "").split(".")
+        return tuple(prefix_parts + ([None] * (4 - len(prefix_parts))))
+
+    def yield_mlmodel_lineage_details(
+        self, mlmodel_details: Any, db_service_prefix: Optional[str] = None
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Get lineage between MLModel and source tables.
+        To be implemented by sources that support lineage.
+        """
+
+    def yield_mlmodel_lineage(
+        self, mlmodel_details: Any
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Yields lineage if config is enabled.
+        We will look for the data in all the services we have informed.
+        """
+        db_service_prefixes = self.get_db_service_prefixes()
+        for db_service_prefix in db_service_prefixes or [None]:
+            for lineage in (
+                self.yield_mlmodel_lineage_details(mlmodel_details, db_service_prefix)
+                or []
+            ):
+                yield from self.yield_lineage_request(lineage)
+
+    def yield_lineage_request(
+        self, lineage: Optional[Either[AddLineageRequest]] = None
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Method to yield lineage request
+        """
+        if lineage:
+            if lineage.right is not None:
+                yield Either(
+                    right=OMetaLineageRequest(
+                        lineage_request=lineage.right,
+                        override_lineage=(
+                            self.source_config.overrideLineage
+                            if hasattr(self.source_config, "overrideLineage")
+                            else False
+                        ),
+                    )
+                )
+            else:
+                yield lineage