openmetadata-ingestion 1.8.0.0__py3-none-any.whl → 1.8.1.0__py3-none-any.whl

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

@@ -86,3 +86,29 @@ DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
 )
 
 DATABRICKS_DDL = "SHOW CREATE TABLE `{table_name}`"
+
+DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB = textwrap.dedent(
+    """
+    SELECT DISTINCT
+        entity_id as job_id,
+        entity_type as job_type,
+        source_table_full_name as source_table_full_name,
+        target_table_full_name as target_table_full_name
+    FROM system.access.table_lineage
+    WHERE entity_type ILIKE 'job'
+    """
+)
+
+DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB = textwrap.dedent(
+    """
+    SELECT DISTINCT
+        entity_id as job_id,
+        entity_type as job_type,
+        source_table_full_name as source_table_full_name,
+        source_column_name as source_column_name,
+        target_table_full_name as target_table_full_name,
+        target_column_name as target_column_name
+    FROM system.access.column_lineage
+    WHERE entity_type ILIKE 'job'
+    """
+)

metadata/ingestion/source/database/postgres/metadata.py

@@ -40,7 +40,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -300,9 +304,15 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         """List stored procedures"""
         if self.source_config.includeStoredProcedures:
             yield from self._get_stored_procedures_internal(
-                POSTGRES_GET_STORED_PROCEDURES
+                POSTGRES_GET_STORED_PROCEDURES.format(
+                    schema_name=self.context.get().database_schema
+                )
+            )
+            yield from self._get_stored_procedures_internal(
+                POSTGRES_GET_FUNCTIONS.format(
+                    schema_name=self.context.get().database_schema
+                )
             )
-            yield from self._get_stored_procedures_internal(POSTGRES_GET_FUNCTIONS)
 
     def yield_stored_procedure(
         self, stored_procedure
@@ -311,7 +321,9 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
                 name=EntityName(stored_procedure.name),
-                description=None,
+                description=Markdown(stored_procedure.description)
+                if stored_procedure.description
+                else None,
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
                     code=stored_procedure.definition,

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

@@ -24,3 +24,4 @@ class PostgresStoredProcedure(BaseModel):
     definition: str
     language: Optional[str] = None
     procedure_type: Optional[str] = Field(None, alias="procedure_type")
+    description: Optional[str] = Field(None, alias="description")

metadata/ingestion/source/database/postgres/queries.py

@@ -225,10 +225,12 @@ POSTGRES_GET_STORED_PROCEDURES = """
         proargtypes AS argument_types,
         prorettype::regtype AS return_type,
         prosrc AS definition,
-        'StoredProcedure' as procedure_type
+        'StoredProcedure' as procedure_type,
+        obj_description(pg_proc.oid, 'pg_proc') AS description
     FROM pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
-    WHERE prokind = 'p';
+    WHERE prokind = 'p'
+    and pg_namespace.nspname = '{schema_name}';
 """
 
 POSTGRES_GET_FUNCTIONS = """
@@ -238,11 +240,12 @@ SELECT
     proargtypes AS argument_types,
     prorettype :: regtype AS return_type,
     prosrc AS definition,
-    'Function' as procedure_type
+    'Function' as procedure_type,
+    obj_description(pg_proc.oid, 'pg_proc') AS description
 FROM
     pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
 WHERE
     prokind = 'f'
-    and pg_namespace.nspname NOT IN ('pg_catalog', 'information_schema');
+    and pg_namespace.nspname = '{schema_name}';
 """

metadata/ingestion/source/pipeline/databrickspipeline/connection.py

@@ -24,17 +24,38 @@ from metadata.generated.schema.entity.services.connections.pipeline.databricksPi
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    init_empty_connection_arguments,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.utils.constants import THREE_MIN
 
 
+def get_connection_url(connection: DatabricksPipelineConnection) -> str:
+    url = f"databricks+connector://token:{connection.token.get_secret_value()}@{connection.hostPort}"
+    return url
+
+
 def get_connection(connection: DatabricksPipelineConnection) -> DatabricksClient:
     """
     Create connection
     """
-    return DatabricksClient(connection)
+
+    if connection.httpPath:
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.root["http_path"] = connection.httpPath
+
+    engine = create_generic_db_connection(
+        connection=connection,
+        get_connection_url_fn=get_connection_url,
+        get_connection_args_fn=get_connection_args_common,
+    )
+    return DatabricksClient(connection, engine)
 
 
 def test_connection(
@@ -49,7 +70,10 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"GetPipelines": client.list_jobs_test_connection}
+    test_fn = {
+        "GetPipelines": client.list_jobs_test_connection,
+        "GetLineage": client.test_lineage_query,
+    }
 
     return test_connection_steps(
         metadata=metadata,

metadata/ingestion/source/pipeline/databrickspipeline/metadata.py

@@ -14,7 +14,7 @@ Databricks pipeline source to extract metadata
 """
 
 import traceback
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Tuple
 
 from pydantic import ValidationError
 
@@ -27,6 +27,7 @@ from metadata.generated.schema.entity.data.pipeline import (
     Task,
     TaskStatus,
 )
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.pipeline.databricksPipelineConnection import (
     DatabricksPipelineConnection,
 )
@@ -43,8 +44,16 @@ from metadata.generated.schema.type.basic import (
     SourceUrl,
     Timestamp,
 )
+from metadata.generated.schema.type.entityLineage import (
+    ColumnLineage,
+    EntitiesEdge,
+    LineageDetails,
+)
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.databrickspipeline.models import (
@@ -54,7 +63,6 @@ from metadata.ingestion.source.pipeline.databrickspipeline.models import (
 from metadata.ingestion.source.pipeline.pipeline_service import PipelineServiceSource
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
 
 logger = ingestion_logger()
 
@@ -200,21 +208,15 @@ class DatabrickspipelineSource(PipelineServiceSource):
                         executionStatus=STATUS_MAP.get(
                             run.state.result_state, StatusType.Failed
                         ),
-                        startTime=Timestamp(
-                            convert_timestamp_to_milliseconds(run.start_time)
-                        ),
-                        endTime=Timestamp(
-                            convert_timestamp_to_milliseconds(run.end_time)
-                        ),
+                        startTime=Timestamp(run.start_time),
+                        endTime=Timestamp(run.end_time) if run.end_time else None,
                         logLink=run.run_page_url,
                     )
                     for task in run.tasks or []
                 ]
                 pipeline_status = PipelineStatus(
                     taskStatus=task_status,
-                    timestamp=Timestamp(
-                        convert_timestamp_to_milliseconds(run.start_time)
-                    ),
+                    timestamp=Timestamp(run.start_time),
                     executionStatus=STATUS_MAP.get(
                         run.state.result_state,
                         StatusType.Failed,
@@ -235,13 +237,173 @@ class DatabrickspipelineSource(PipelineServiceSource):
         except Exception as exc:
             yield Either(
                 left=StackTraceError(
-                    name=pipeline_fqn,
+                    name=pipeline_details.job_id,
                     error=f"Failed to yield pipeline status: {exc}",
                     stackTrace=traceback.format_exc(),
                 )
             )
 
+    def _process_and_validate_column_lineage(
+        self,
+        column_lineage: List[Tuple[str, str]],
+        from_entity: Table,
+        to_entity: Table,
+    ) -> List[ColumnLineage]:
+        """
+        Process and validate column lineage
+        """
+        processed_column_lineage = []
+        if column_lineage:
+            for column_tuple in column_lineage or []:
+                try:
+                    if len(column_tuple) < 2:
+                        logger.debug(f"Skipping invalid column tuple: {column_tuple}")
+                        continue
+
+                    source_col = column_tuple[0]
+                    target_col = column_tuple[-1]
+
+                    if not source_col or not target_col:
+                        logger.debug(
+                            f"Skipping column tuple with empty values: source={source_col}, "
+                            f"target={target_col}, to_entity={to_entity.name}"
+                        )
+                        continue
+
+                    from_column = get_column_fqn(
+                        table_entity=from_entity, column=str(source_col)
+                    )
+                    to_column = get_column_fqn(
+                        table_entity=to_entity,
+                        column=str(target_col),
+                    )
+                    if from_column and to_column:
+                        processed_column_lineage.append(
+                            ColumnLineage(
+                                fromColumns=[from_column],
+                                toColumn=to_column,
+                            )
+                        )
+                except Exception as err:
+                    logger.warning(
+                        f"Error processing column lineage {column_tuple}: {err}"
+                    )
+                    logger.debug(traceback.format_exc())
+                    continue
+        if not processed_column_lineage:
+            logger.warning(
+                f"No column lineage found for {from_entity.name} to {to_entity.name}"
+            )
+        return processed_column_lineage or []
+
     def yield_pipeline_lineage_details(
         self, pipeline_details: DataBrickPipelineDetails
     ) -> Iterable[Either[AddLineageRequest]]:
-        """Get lineage between pipeline and data sources. Not Implemented."""
+        try:
+            pipeline_fqn = fqn.build(
+                metadata=self.metadata,
+                entity_type=Pipeline,
+                service_name=self.context.get().pipeline_service,
+                pipeline_name=self.context.get().pipeline,
+            )
+
+            pipeline_entity = self.metadata.get_by_name(
+                entity=Pipeline, fqn=pipeline_fqn
+            )
+
+            table_lineage_list = self.client.get_table_lineage(
+                job_id=pipeline_details.job_id
+            )
+            if table_lineage_list:
+                for table_lineage in table_lineage_list:
+                    source_table_full_name = table_lineage.get("source_table_full_name")
+                    target_table_full_name = table_lineage.get("target_table_full_name")
+                    if source_table_full_name and target_table_full_name:
+                        source = fqn.split_table_name(source_table_full_name)
+                        target = fqn.split_table_name(target_table_full_name)
+                        for dbservicename in self.get_db_service_names() or ["*"]:
+
+                            from_entity = self.metadata.get_by_name(
+                                entity=Table,
+                                fqn=fqn.build(
+                                    metadata=self.metadata,
+                                    entity_type=Table,
+                                    table_name=source.get("table"),
+                                    database_name=source.get("database"),
+                                    schema_name=source.get("database_schema"),
+                                    service_name=dbservicename,
+                                ),
+                            )
+
+                            if from_entity is None:
+                                continue
+
+                            to_entity = self.metadata.get_by_name(
+                                entity=Table,
+                                fqn=fqn.build(
+                                    metadata=self.metadata,
+                                    entity_type=Table,
+                                    table_name=target.get("table"),
+                                    database_name=target.get("database"),
+                                    schema_name=target.get("database_schema"),
+                                    service_name=dbservicename,
+                                ),
+                            )
+
+                            if to_entity is None:
+                                continue
+
+                            processed_column_lineage = (
+                                self._process_and_validate_column_lineage(
+                                    column_lineage=self.client.get_column_lineage(
+                                        job_id=pipeline_details.job_id,
+                                        TableKey=(
+                                            source_table_full_name,
+                                            target_table_full_name,
+                                        ),
+                                    ),
+                                    from_entity=from_entity,
+                                    to_entity=to_entity,
+                                )
+                            )
+
+                            lineage_details = LineageDetails(
+                                pipeline=EntityReference(
+                                    id=pipeline_entity.id.root, type="pipeline"
+                                ),
+                                source=LineageSource.PipelineLineage,
+                                columnsLineage=processed_column_lineage,
+                            )
+
+                            yield Either(
+                                right=AddLineageRequest(
+                                    edge=EntitiesEdge(
+                                        fromEntity=EntityReference(
+                                            id=from_entity.id,
+                                            type="table",
+                                        ),
+                                        toEntity=EntityReference(
+                                            id=to_entity.id,
+                                            type="table",
+                                        ),
+                                        lineageDetails=lineage_details,
+                                    )
+                                )
+                            )
+
+                else:
+                    logger.debug(
+                        f"No source or target table full name found for job {pipeline_details.job_id}"
+                    )
+            else:
+                logger.debug(
+                    f"No table lineage found for job {pipeline_details.job_id}"
+                )
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=pipeline_details.job_id,
+                    error=f"Wild error ingesting pipeline lineage {pipeline_details} - {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )

metadata/pii/algorithms/classifiers.py

@@ -103,7 +103,6 @@ class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
         column_name: Optional[str] = None,
         column_data_type: Optional[DataType] = None,
     ) -> Mapping[PIITag, float]:
-
         if column_data_type is not None and is_non_pii_datatype(column_data_type):
             return {}
 

metadata/pii/algorithms/presidio_patches.py

@@ -16,6 +16,10 @@ from typing import List, Protocol, Sequence
 from dateutil.parser import parse
 from presidio_analyzer import RecognizerResult
 
+from metadata.utils.logger import pii_logger
+
+logger = pii_logger()
+
 
 class PresidioRecognizerResultPatcher(Protocol):
     """
@@ -76,8 +80,11 @@ def date_time_patcher(
             # try to parse using dateutils, if it fails, skip the result
             try:
                 _ = parse(text[result.start : result.end])
-            except ValueError:
+            except (ValueError, OverflowError):
                 # if parsing fails, skip the result
                 continue
+            except Exception as e:
+                logger.info("Unexpected error while parsing date time: %s", e)
+                continue
         patched_result.append(result)
     return patched_result

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -18,9 +18,13 @@ from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class DatabricksSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     def get_client(self):
         """client is the session for SQA"""
-        self.connection = databricks_get_connection(self.service_connection_config)
         client = super().get_client()
         self.set_catalog(client)
         return client

metadata/sampler/sqlalchemy/sampler.py

@@ -13,6 +13,7 @@ Helper module to handle data sampling
 for the profiler
 """
 import hashlib
+from contextlib import contextmanager
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
@@ -71,6 +72,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         self._table = self.build_table_orm(
             self.entity, self.service_connection_config, self.ometa_client
         )
+        self._active_sessions = set()
 
     @property
     def raw_dataset(self):
@@ -79,7 +81,20 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
     def get_client(self):
         """Build the SQA Client"""
         session_factory = create_and_bind_thread_safe_session(self.connection)
-        return session_factory()
+        session = session_factory()
+        self._active_sessions.add(session)
+        return session
+
+    @contextmanager
+    def get_client_context(self):
+        """Get client as context manager for proper cleanup"""
+        session = self.get_client()
+        try:
+            yield session
+        finally:
+            if session in self._active_sessions:
+                self._active_sessions.remove(session)
+            session.close()
 
     def set_tablesample(self, selectable: Table):
         """Set the tablesample for the table. To be implemented by the child SQA sampler class
@@ -199,7 +214,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
             sqa_sample = (
                 client.query(*sqa_columns)
                 .select_from(ds)
@@ -217,7 +232,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not is_safe_sql_query(self.sample_query):
             raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
             rnd = client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
@@ -264,5 +279,6 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def close(self):
         """Close the connection"""
-        self.get_client().close()
+        for session in self._active_sessions:
+            session.close()
         self.connection.pool.dispose()

metadata/sampler/sqlalchemy/unitycatalog/sampler.py

@@ -20,9 +20,13 @@ from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class UnityCatalogSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     def get_client(self):
         """client is the session for SQA"""
-        self.connection = databricks_get_connection(self.service_connection_config)
         client = super().get_client()
         self.set_catalog(client)
         return client