openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/ingestion/source/database/database_service.py

@@ -11,6 +11,7 @@
 """
 Base class for ingesting database services
 """
+import traceback
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Set, Tuple, Union
 
@@ -50,6 +51,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
@@ -60,8 +62,8 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.source.connections import get_test_connection_fn
 from metadata.utils import fqn
@@ -111,6 +113,13 @@ class DatabaseServiceTopology(ServiceTopology):
     database = TopologyNode(
         producer="get_database_names",
         stages=[
+            NodeStage(
+                type_=OMetaTagAndClassification,
+                context="tags",
+                processor="yield_database_tag_details",
+                nullable=True,
+                store_all_in_context=True,
+            ),
             NodeStage(
                 type_=Database,
                 context="database",
@@ -118,7 +127,7 @@ class DatabaseServiceTopology(ServiceTopology):
                 consumer=["database_service"],
                 cache_entities=True,
                 use_cache=True,
-            )
+            ),
         ],
         children=["databaseSchema"],
     )
@@ -130,7 +139,7 @@ class DatabaseServiceTopology(ServiceTopology):
                 context="tags",
                 processor="yield_database_schema_tag_details",
                 nullable=True,
-                cache_all=True,
+                store_all_in_context=True,
             ),
             NodeStage(
                 type_=DatabaseSchema,
@@ -152,7 +161,7 @@ class DatabaseServiceTopology(ServiceTopology):
                 context="tags",
                 processor="yield_table_tag_details",
                 nullable=True,
-                cache_all=True,
+                store_all_in_context=True,
             ),
             NodeStage(
                 type_=Table,
@@ -176,7 +185,8 @@ class DatabaseServiceTopology(ServiceTopology):
                 context="stored_procedures",
                 processor="yield_stored_procedure",
                 consumer=["database_service", "database", "database_schema"],
-                cache_all=True,
+                store_all_in_context=True,
+                store_fqn=True,
                 use_cache=True,
             ),
         ],
@@ -202,7 +212,11 @@ class DatabaseServiceSource(
     inspector: Inspector
 
     topology = DatabaseServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
+
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
 
     def prepare(self):
         """By default, there is no preparation needed"""
@@ -270,9 +284,16 @@ class DatabaseServiceSource(
         From topology. To be run for each schema
         """
 
+    def yield_database_tag(
+        self, database_name: str
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each database
+        """
+
     def yield_table_tags(
         self, table_name_and_type: Tuple[str, TableType]
-    ) -> Iterable[Either[CreateTableRequest]]:
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
         """
         From topology. To be run for each table
         """
@@ -295,6 +316,15 @@ class DatabaseServiceSource(
         if self.source_config.includeTags:
             yield from self.yield_tag(schema_name) or []
 
+    def yield_database_tag_details(
+        self, database_name: str
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each database
+        """
+        if self.source_config.includeTags:
+            yield from self.yield_database_tag(database_name) or []
+
     @abstractmethod
     def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
         """
@@ -361,6 +391,35 @@ class DatabaseServiceSource(
                     tag_labels.append(tag_label)
         return tag_labels or None
 
+    def get_database_tag_labels(self, database_name: str) -> Optional[List[TagLabel]]:
+        """
+        Method to get schema tags
+        This will only get executed if the tags context
+        is properly informed
+        """
+        database_fqn = fqn.build(
+            self.metadata,
+            entity_type=Database,
+            service_name=self.context.database_service,
+            database_name=database_name,
+        )
+        return self.get_tag_by_fqn(entity_fqn=database_fqn)
+
+    def get_schema_tag_labels(self, schema_name: str) -> Optional[List[TagLabel]]:
+        """
+        Method to get schema tags
+        This will only get executed if the tags context
+        is properly informed
+        """
+        schema_fqn = fqn.build(
+            self.metadata,
+            entity_type=DatabaseSchema,
+            service_name=self.context.database_service,
+            database_name=self.context.database,
+            schema_name=schema_name,
+        )
+        return self.get_tag_by_fqn(entity_fqn=schema_fqn)
+
     def get_tag_labels(self, table_name: str) -> Optional[List[TagLabel]]:
         """
         This will only get executed if the tags context
@@ -448,10 +507,33 @@ class DatabaseServiceSource(
                 continue
             yield schema_fqn if return_fqn else schema_name
 
+    def get_owner_ref(self, table_name: str) -> Optional[EntityReference]:
+        """
+        Method to process the table owners
+        """
+        try:
+            if self.source_config.includeOwners:
+                owner_name = self.inspector.get_table_owner(
+                    connection=self.connection,  # pylint: disable=no-member
+                    table_name=table_name,
+                    schema=self.context.database_schema,
+                )
+                owner_ref = self.metadata.get_reference_by_name(name=owner_name)
+                return owner_ref
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error processing owner for table {table_name}: {exc}")
+        return None
+
     def mark_tables_as_deleted(self):
         """
         Use the current inspector to mark tables as deleted
         """
+        if not self.context.__dict__.get("database"):
+            raise ValueError(
+                "No Database found in the context. We cannot run the table deletion."
+            )
+
         if self.source_config.markDeletedTables:
             logger.info(
                 f"Mark Deleted Tables set to True. Processing database [{self.context.database}]"

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

@@ -14,7 +14,7 @@ Client to interact with databricks apis
 import json
 import traceback
 from datetime import timedelta
-from typing import List
+from typing import Iterable, List
 
 import requests
 
@@ -22,19 +22,20 @@ from metadata.generated.schema.entity.services.connections.database.databricksCo
     DatabricksConnection,
 )
 from metadata.ingestion.ometa.client import APIError
-from metadata.ingestion.source.database.databricks.models import (
-    LineageColumnStreams,
-    LineageTableStreams,
-)
 from metadata.utils.constants import QUERY_WITH_DBT, QUERY_WITH_OM_VERSION
 from metadata.utils.helpers import datetime_to_ts
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 API_TIMEOUT = 10
+PAGE_SIZE = 100
 QUERIES_PATH = "/sql/history/queries"
-TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage/get"
-COLUMN_LINEAGE_PATH = "/lineage-tracking/column-lineage/get"
+
+
+class DatabricksClientException(Exception):
+    """
+    Class to throw auth and other databricks api exceptions.
+    """
 
 
 class DatabricksClient:
@@ -66,15 +67,33 @@ class DatabricksClient:
         if res.status_code != 200:
             raise APIError(res.json)
 
+    def _run_query_paginator(self, data, result, end_time, response):
+        while True:
+            if response:
+                next_page_token = response.get("next_page_token", None)
+                has_next_page = response.get("has_next_page", None)
+                if next_page_token:
+                    data["page_token"] = next_page_token
+                if not has_next_page:
+                    data = {}
+                    break
+            else:
+                break
+
+            if result[-1]["execution_end_time_ms"] <= end_time:
+                response = self.client.get(
+                    self.base_query_url,
+                    data=json.dumps(data),
+                    headers=self.headers,
+                    timeout=API_TIMEOUT,
+                ).json()
+                yield from response.get("res") or []
+
     def list_query_history(self, start_date=None, end_date=None) -> List[dict]:
         """
         Method returns List the history of queries through SQL warehouses
         """
-        query_details = []
         try:
-            next_page_token = None
-            has_next_page = None
-
             data = {}
             daydiff = end_date - start_date
 
@@ -104,36 +123,15 @@ class DatabricksClient:
                     result = response.get("res") or []
                     data = {}
 
-                while True:
-                    if result:
-                        query_details.extend(result)
-
-                        next_page_token = response.get("next_page_token", None)
-                        has_next_page = response.get("has_next_page", None)
-                        if next_page_token:
-                            data["page_token"] = next_page_token
-
-                        if not has_next_page:
-                            data = {}
-                            break
-                    else:
-                        break
-
-                    if result[-1]["execution_end_time_ms"] <= end_time:
-                        response = self.client.get(
-                            self.base_query_url,
-                            data=json.dumps(data),
-                            headers=self.headers,
-                            timeout=API_TIMEOUT,
-                        ).json()
-                        result = response.get("res")
+                yield from result
+                yield from self._run_query_paginator(
+                    data=data, result=result, end_time=end_time, response=response
+                ) or []
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(exc)
 
-        return query_details
-
     def is_query_valid(self, row) -> bool:
         query_text = row.get("query_text")
         return not (
@@ -141,13 +139,24 @@ class DatabricksClient:
             or query_text.startswith(QUERY_WITH_OM_VERSION)
         )
 
-    def list_jobs(self) -> List[dict]:
+    def list_jobs_test_connection(self) -> None:
+        data = {"limit": 1, "expand_tasks": True, "offset": 0}
+        response = self.client.get(
+            self.jobs_list_url,
+            data=json.dumps(data),
+            headers=self.headers,
+            timeout=API_TIMEOUT,
+        )
+        if response.status_code != 200:
+            raise DatabricksClientException(response.text)
+
+    def list_jobs(self) -> Iterable[dict]:
         """
         Method returns List all the created jobs in a Databricks Workspace
         """
-        job_list = []
         try:
-            data = {"limit": 25, "expand_tasks": True, "offset": 0}
+            iteration_count = 1
+            data = {"limit": PAGE_SIZE, "expand_tasks": True, "offset": 0}
 
             response = self.client.get(
                 self.jobs_list_url,
@@ -156,10 +165,10 @@ class DatabricksClient:
                 timeout=API_TIMEOUT,
             ).json()
 
-            job_list.extend(response.get("jobs") or [])
+            yield from response.get("jobs") or []
 
-            while response["has_more"]:
-                data["offset"] = len(response.get("jobs") or [])
+            while response and response.get("has_more"):
+                data["offset"] = PAGE_SIZE * iteration_count
 
                 response = self.client.get(
                     self.jobs_list_url,
@@ -167,20 +176,17 @@ class DatabricksClient:
                     headers=self.headers,
                     timeout=API_TIMEOUT,
                 ).json()
-
-                job_list.extend(response.get("jobs") or [])
+                iteration_count += 1
+                yield from response.get("jobs") or []
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(exc)
 
-        return job_list
-
     def get_job_runs(self, job_id) -> List[dict]:
         """
         Method returns List of all runs for a job by the specified job_id
         """
-        job_runs = []
         try:
             params = {
                 "job_id": job_id,
@@ -197,7 +203,7 @@ class DatabricksClient:
                 timeout=API_TIMEOUT,
             ).json()
 
-            job_runs.extend(response.get("runs") or [])
+            yield from response.get("runs") or []
 
             while response["has_more"]:
                 params.update({"start_time_to": response["runs"][-1]["start_time"]})
@@ -209,62 +215,8 @@ class DatabricksClient:
                     timeout=API_TIMEOUT,
                 ).json()
 
-                job_runs.extend(response.get("runs" or []))
+                yield from response.get("runs") or []
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(exc)
-
-        return job_runs
-
-    def get_table_lineage(self, table_name: str) -> LineageTableStreams:
-        """
-        Method returns table lineage details
-        """
-        try:
-            data = {
-                "table_name": table_name,
-            }
-
-            response = self.client.get(
-                f"{self.base_url}{TABLE_LINEAGE_PATH}",
-                headers=self.headers,
-                data=json.dumps(data),
-                timeout=API_TIMEOUT,
-            ).json()
-            if response:
-                return LineageTableStreams(**response)
-
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.error(exc)
-
-        return LineageTableStreams()
-
-    def get_column_lineage(
-        self, table_name: str, column_name: str
-    ) -> LineageColumnStreams:
-        """
-        Method returns table lineage details
-        """
-        try:
-            data = {
-                "table_name": table_name,
-                "column_name": column_name,
-            }
-
-            response = self.client.get(
-                f"{self.base_url}{COLUMN_LINEAGE_PATH}",
-                headers=self.headers,
-                data=json.dumps(data),
-                timeout=API_TIMEOUT,
-            ).json()
-
-            if response:
-                return LineageColumnStreams(**response)
-
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.error(exc)
-
-        return LineageColumnStreams()

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

@@ -13,9 +13,8 @@
 Source connection handler
 """
 from functools import partial
-from typing import Optional, Union
+from typing import Optional
 
-from databricks.sdk import WorkspaceClient
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import DatabaseError
 from sqlalchemy.inspection import inspect
@@ -37,11 +36,9 @@ from metadata.ingestion.connections.test_connections import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.databricks.client import DatabricksClient
-from metadata.ingestion.source.database.databricks.models import DatabricksTable
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
 )
-from metadata.utils.db_utils import get_host_from_host_port
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -52,17 +49,11 @@ def get_connection_url(connection: DatabricksConnection) -> str:
     return url
 
 
-def get_connection(connection: DatabricksConnection) -> Union[Engine, WorkspaceClient]:
+def get_connection(connection: DatabricksConnection) -> Engine:
     """
     Create connection
     """
 
-    if connection.useUnityCatalog:
-        return WorkspaceClient(
-            host=get_host_from_host_port(connection.hostPort),
-            token=connection.token.get_secret_value(),
-        )
-
     if connection.httpPath:
         if not connection.connectionArguments:
             connection.connectionArguments = init_empty_connection_arguments()
@@ -77,7 +68,7 @@ def get_connection(connection: DatabricksConnection) -> Union[Engine, WorkspaceC
 
 def test_connection(
     metadata: OpenMetadata,
-    connection: Union[Engine, WorkspaceClient],
+    connection: Engine,
     service_connection: DatabricksConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
 ) -> None:
@@ -99,49 +90,19 @@ def test_connection(
         except DatabaseError as soe:
             logger.debug(f"Failed to fetch catalogs due to: {soe}")
 
-    if service_connection.useUnityCatalog:
-        table_obj = DatabricksTable()
-
-        def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
-            for catalog in connection.catalogs.list():
-                table_obj.catalog_name = catalog.name
-                break
-
-        def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
-            for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
-                table_obj.schema_name = schema.name
-                break
-
-        def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
-            for table in connection.tables.list(
-                catalog_name=table_obj.catalog_name, schema_name=table_obj.schema_name
-            ):
-                table_obj.name = table.name
-                break
-
-        test_fn = {
-            "CheckAccess": connection.catalogs.list,
-            "GetDatabases": partial(get_catalogs, connection, table_obj),
-            "GetSchemas": partial(get_schemas, connection, table_obj),
-            "GetTables": partial(get_tables, connection, table_obj),
-            "GetViews": partial(get_tables, connection, table_obj),
-            "GetQueries": client.test_query_api_access,
-        }
-
-    else:
-        inspector = inspect(connection)
-        test_fn = {
-            "CheckAccess": partial(test_connection_engine_step, connection),
-            "GetSchemas": inspector.get_schema_names,
-            "GetTables": inspector.get_table_names,
-            "GetViews": inspector.get_view_names,
-            "GetDatabases": partial(
-                test_database_query,
-                engine=connection,
-                statement=DATABRICKS_GET_CATALOGS,
-            ),
-            "GetQueries": client.test_query_api_access,
-        }
+    inspector = inspect(connection)
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, connection),
+        "GetSchemas": inspector.get_schema_names,
+        "GetTables": inspector.get_table_names,
+        "GetViews": inspector.get_view_names,
+        "GetDatabases": partial(
+            test_database_query,
+            engine=connection,
+            statement=DATABRICKS_GET_CATALOGS,
+        ),
+        "GetQueries": client.test_query_api_access,
+    }
 
     test_connection_steps(
         metadata=metadata,

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

@@ -11,38 +11,41 @@
 """
 Databricks lineage module
 """
-from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
-    DatabricksConnection,
-)
-from metadata.generated.schema.metadataIngestion.workflow import (
-    Source as WorkflowSource,
-)
-from metadata.ingestion.api.steps import InvalidSourceException
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.databricks.legacy.lineage import (
-    DatabricksLineageLegacySource,
-)
-from metadata.ingestion.source.database.databricks.unity_catalog.lineage import (
-    DatabricksUnityCatalogLineageSource,
+import traceback
+from datetime import datetime
+from typing import Iterator
+
+from metadata.generated.schema.type.tableQuery import TableQuery
+from metadata.ingestion.source.database.databricks.query_parser import (
+    DatabricksQueryParserSource,
 )
+from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class DatabricksLineageSource:
+class DatabricksLineageSource(DatabricksQueryParserSource, LineageSource):
     """
-    Databricks Lineage Source
+    Databricks Lineage Legacy Source
     """
 
-    @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatabricksConnection = config.serviceConnection.__root__.config
-        if not isinstance(connection, DatabricksConnection):
-            raise InvalidSourceException(
-                f"Expected DatabricksConnection, but got {connection}"
-            )
-        if not connection.useUnityCatalog:
-            return DatabricksLineageLegacySource(config, metadata)
-        return DatabricksUnityCatalogLineageSource(config, metadata)
+    def yield_table_query(self) -> Iterator[TableQuery]:
+        data = self.client.list_query_history(
+            start_date=self.start,
+            end_date=self.end,
+        )
+        for row in data or []:
+            try:
+                if self.client.is_query_valid(row):
+                    yield TableQuery(
+                        query=row.get("query_text"),
+                        userName=row.get("user_name"),
+                        startTime=row.get("query_start_time_ms"),
+                        endTime=row.get("execution_end_time_ms"),
+                        analysisDate=datetime.now(),
+                        serviceName=self.config.serviceName,
+                    )
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error processing query_dict {row}: {exc}")