openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

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(
@@ -111,16 +143,15 @@ def test_connection(
 
     def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
         for catalog in connection.catalogs.list():
-            table_obj.catalog_name = catalog.name
-            break
+            if catalog.name != "__databricks_internal":
+                table_obj.catalog_name = catalog.name
+                return
 
     def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
-        for catalog in connection.catalogs.list():
-            for schema in connection.schemas.list(catalog_name=catalog.name):
-                if schema.name:
-                    table_obj.schema_name = schema.name
-                    table_obj.catalog_name = catalog.name
-                    return
+        for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
+            if schema.name:
+                table_obj.schema_name = schema.name
+                return
 
     def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
         if table_obj.catalog_name and table_obj.schema_name:

metadata/ingestion/source/database/unitycatalog/data_diff/__init__.py

@@ -0,0 +1 @@
+"""Unity Catalog data diff module"""

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

@@ -0,0 +1,12 @@
+"""Unity Catalog spec for data diff - delegates to Databricks connection"""
+
+from metadata.ingestion.source.database.common.data_diff.databricks_base import (
+    DatabricksBaseTableParameter,
+)
+
+
+class UnityCatalogTableParameter(DatabricksBaseTableParameter):
+    """Unity Catalog table parameter setter - uses Unity Catalog connection
+    which is databricks-based for data diff operations"""
+
+    pass

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

@@ -15,6 +15,7 @@ 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
@@ -90,32 +91,138 @@ 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]:
-        col_lineage = []
-        for column in to_table.columns:
-            column_streams = self.client.get_column_lineage(
-                databricks_table_fqn, column_name=column.name.root
+        try:
+            col_lineage = []
+            for column in to_table.columns:
+                column_streams = self.client.get_column_lineage(
+                    databricks_table_fqn, column_name=column.name.root
+                )
+                from_columns = []
+                for col in column_streams.upstream_cols:
+                    col_fqn = get_column_fqn(from_table, col.name)
+                    if col_fqn:
+                        from_columns.append(col_fqn)
+
+                if from_columns:
+                    col_lineage.append(
+                        ColumnLineage(
+                            fromColumns=from_columns,
+                            toColumn=column.fullyQualifiedName.root,
+                        )
+                    )
+            if col_lineage:
+                return LineageDetails(
+                    columnsLineage=col_lineage, source=LineageSource.QueryLineage
+                )
+            return None
+        except Exception as exc:
+            logger.debug(
+                f"Error computing column lineage for {to_table.fullyQualifiedName.root} - {exc}"
+            )
+            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
+
+            storage_location = file_info.storage_location.rstrip("/")
+            location_entity = self.metadata.es_search_container_by_path(
+                full_path=storage_location, fields="dataModel"
             )
-            from_columns = []
-            for col in column_streams.upstream_cols:
-                col_fqn = get_column_fqn(from_table, col.name)
-                if col_fqn:
-                    from_columns.append(col_fqn)
-
-            if from_columns:
-                col_lineage.append(
-                    ColumnLineage(
-                        fromColumns=from_columns,
-                        toColumn=column.fullyQualifiedName.root,
+
+            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: {storage_location}"
                 )
-        if col_lineage:
-            return LineageDetails(
-                columnsLineage=col_lineage, source=LineageSource.QueryLineage
+        except Exception as exc:
+            logger.debug(
+                f"Error while processing external location lineage for {file_info.storage_location}: {exc}"
             )
-        return None
+            logger.debug(traceback.format_exc())
 
     def _handle_upstream_table(
         self,
@@ -123,10 +230,18 @@ class UnitycatalogLineageSource(Source):
         table: Table,
         databricks_table_fqn: str,
     ) -> Iterable[Either[AddLineageRequest]]:
-        for upstream_table in table_streams.upstream_tables:
+        for upstream_entity in table_streams.upstreams:
             try:
-                if not upstream_table.name:
+                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:
                     continue
+
+                upstream_table = upstream_entity.tableInfo
                 from_entity_fqn = fqn.build(
                     metadata=self.metadata,
                     entity_type=Table,
@@ -157,11 +272,77 @@ class UnitycatalogLineageSource(Source):
                             )
                         ),
                     )
+                else:
+                    logger.debug(
+                        f"Unable to find upstream entity for "
+                        f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
+                        f" -> {databricks_table_fqn}"
+                    )
             except Exception:
                 logger.debug(
-                    "Error while processing lineage for "
-                    f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
-                    f" -> {databricks_table_fqn}"
+                    "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}"
                 )
                 logger.debug(traceback.format_exc())
 
@@ -212,10 +393,18 @@ 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
                     )
 
+                    # Disabling downstream lineage for now as it causes slowness
+                    # 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/metadata.py

@@ -92,7 +92,7 @@ logger = ingestion_logger()
 UNITY_CATALOG_TAG = "UNITY CATALOG TAG"
 UNITY_CATALOG_TAG_CLASSIFICATION = "UNITY CATALOG TAG CLASSIFICATION"
 
-
+# pylint: disable=protected-access
 class UnitycatalogSource(
     ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
 ):
@@ -306,6 +306,8 @@ class UnitycatalogSource(
                 if table.table_type:
                     if table.table_type.value.lower() == TableType.View.value.lower():
                         table_type: TableType = TableType.View
+                    if table.table_type.value.lower() == "materialized_view":
+                        table_type: TableType = TableType.MaterializedView
                     elif (
                         table.table_type.value.lower()
                         == TableType.External.value.lower()
@@ -427,7 +429,6 @@ class UnitycatalogSource(
             )
             if referred_table_fqn:
                 for parent_column in column.parent_columns:
-                    # pylint: disable=protected-access
                     col_fqn = fqn._build(referred_table_fqn, parent_column, quote=False)
                     if col_fqn:
                         referred_column_fqns.append(FullyQualifiedEntityName(col_fqn))
@@ -563,7 +564,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,
@@ -617,7 +618,7 @@ class UnitycatalogSource(
                             yield from get_ometa_tag_and_classification(
                                 tag_fqn=FullyQualifiedEntityName(
                                     fqn._build(*tag_fqn_builder(tag))
-                                ),  # pylint: disable=protected-access
+                                ),
                                 tags=[tag.tag_value],
                                 classification_name=tag.tag_name,
                                 tag_description=UNITY_CATALOG_TAG,

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

@@ -22,6 +22,8 @@ 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):
@@ -31,9 +33,23 @@ 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):
-    upstream_tables: Optional[List[DatabricksTable]] = []
-    downstream_tables: Optional[List[DatabricksTable]] = []
+    upstreams: Optional[List[LineageEntity]] = []
+    downstreams: Optional[List[LineageEntity]] = []
 
 
 class LineageColumnStreams(BaseModel):

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

@@ -1,6 +1,9 @@
 from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
     UnityCatalogTestSuiteInterface,
 )
+from metadata.ingestion.source.database.unitycatalog.data_diff.data_diff import (
+    UnityCatalogTableParameter,
+)
 from metadata.ingestion.source.database.unitycatalog.lineage import (
     UnitycatalogLineageSource,
 )
@@ -23,4 +26,5 @@ ServiceSpec = DefaultDatabaseSpec(
     profiler_class=UnityCatalogProfilerInterface,
     test_suite_class=UnityCatalogTestSuiteInterface,
     sampler_class=UnityCatalogSamplerInterface,
+    data_diff=UnityCatalogTableParameter,
 )

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: