openmetadata-ingestion 1.2.5.1__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

metadata/ingestion/source/database/common_db_source.py

@@ -41,13 +41,16 @@ from metadata.generated.schema.entity.data.table import (
     TablePartition,
     TableType,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -171,6 +174,7 @@ class CommonDbSourceService(
                 service=self.context.database_service,
                 description=self.get_database_description(database_name),
                 sourceUrl=self.get_source_url(database_name=database_name),
+                tags=self.get_database_tag_labels(database_name=database_name),
             )
         )
 
@@ -209,6 +213,7 @@ class CommonDbSourceService(
                     database_name=self.context.database,
                     schema_name=schema_name,
                 ),
+                tags=self.get_schema_tag_labels(schema_name=schema_name),
             )
         )
 
@@ -458,6 +463,7 @@ class CommonDbSourceService(
                     database_name=self.context.database,
                     table_type=table_type,
                 ),
+                owner=self.get_owner_ref(table_name=table_name),
             )
 
             is_partitioned, partition_details = self.get_table_partition_details(
@@ -488,7 +494,7 @@ class CommonDbSourceService(
             error = f"Unexpected exception to yield table [{table_name}]: {exc}"
             yield Either(
                 left=StackTraceError(
-                    name=table_name, error=error, stack_trace=traceback.format_exc()
+                    name=table_name, error=error, stackTrace=traceback.format_exc()
                 )
             )
 

metadata/ingestion/source/database/common_nosql_source.py

@@ -29,13 +29,16 @@ from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 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, TableType
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
@@ -43,7 +46,7 @@ from metadata.ingestion.source.database.database_service import DatabaseServiceS
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.utils import fqn
 from metadata.utils.constants import DEFAULT_DATABASE
-from metadata.utils.datalake.datalake_utils import get_columns
+from metadata.utils.datalake.datalake_utils import DataFrameColumnParser
 from metadata.utils.filters import filter_by_schema, filter_by_table
 from metadata.utils.logger import ingestion_logger
 
@@ -214,7 +217,8 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         try:
             data = self.get_table_columns_dict(schema_name, table_name)
             df = pd.DataFrame.from_records(list(data))
-            columns = get_columns(df)
+            column_parser = DataFrameColumnParser.create(df)
+            columns = column_parser.get_columns()
             table_request = CreateTableRequest(
                 name=table_name,
                 tableType=table_type,
@@ -242,7 +246,7 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception to yield table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -22,19 +22,14 @@ 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):
@@ -99,7 +94,6 @@ class DatabricksClient:
         Method returns List the history of queries through SQL warehouses
         """
         try:
-
             data = {}
             daydiff = end_date - start_date
 
@@ -161,7 +155,8 @@ class DatabricksClient:
         Method returns List all the created jobs in a Databricks Workspace
         """
         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,
@@ -173,7 +168,7 @@ class DatabricksClient:
             yield from response.get("jobs") or []
 
             while response and response.get("has_more"):
-                data["offset"] = len(response.get("jobs") or [])
+                data["offset"] = PAGE_SIZE * iteration_count
 
                 response = self.client.get(
                     self.jobs_list_url,
@@ -181,7 +176,7 @@ class DatabricksClient:
                     headers=self.headers,
                     timeout=API_TIMEOUT,
                 ).json()
-
+                iteration_count += 1
                 yield from response.get("jobs") or []
 
         except Exception as exc:

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}")