openmetadata-ingestion 1.8.2.0__py3-none-any.whl → 1.8.3.1__py3-none-any.whl

metadata/great_expectations/action.py

@@ -21,6 +21,7 @@ from datetime import datetime
 from typing import Dict, List, Optional, Union, cast
 
 from great_expectations.checkpoint.actions import ValidationAction
+from great_expectations.core import ExpectationConfiguration
 from great_expectations.core.batch import Batch
 from great_expectations.core.batch_spec import (
     RuntimeDataBatchSpec,
@@ -111,6 +112,7 @@ class OpenMetadataValidationAction(ValidationAction):
         self.schema_name = schema_name  # for database without schema concept
         self.config_file_path = config_file_path
         self.ometa_conn = self._create_ometa_connection()
+        self.expectation_suite = None
 
     def _run(  # pylint: disable=unused-argument
         self,
@@ -133,6 +135,12 @@ class OpenMetadataValidationAction(ValidationAction):
             checkpoint_identifier: identifier for the checkpoint
         """
 
+        if expectation_suite_identifier:
+            expectation_suite_name = expectation_suite_identifier.expectation_suite_name
+            self.expectation_suite = self.data_context.get_expectation_suite(
+                expectation_suite_name
+            )
+
         check_point_spec = self._get_checkpoint_batch_spec(data_asset)
         table_entity = None
         if isinstance(check_point_spec, SqlAlchemyDatasourceBatchSpec):
@@ -316,6 +324,18 @@ class OpenMetadataValidationAction(ValidationAction):
 
     def _get_test_case_params_value(self, result: dict) -> List[TestCaseParameterValue]:
         """Build test case parameter value from GE test result"""
+        if self.expectation_suite:
+            expectation = self._get_expectation_config(result)
+            if expectation:
+                return [
+                    TestCaseParameterValue(
+                        name=key,
+                        value=str(value),
+                    )  # type: ignore
+                    for key, value in expectation.kwargs.items()  # type: ignore
+                    if key not in {"column", "batch_id"}
+                ]
+
         if "observed_value" not in result["result"]:
             return [
                 TestCaseParameterValue(
@@ -337,6 +357,17 @@ class OpenMetadataValidationAction(ValidationAction):
         self, result: dict
     ) -> List[TestCaseParameterDefinition]:
         """Build test case parameter definition from GE test result"""
+        if self.expectation_suite:
+            expectation = self._get_expectation_config(result)
+            if expectation:
+                return [
+                    TestCaseParameterDefinition(
+                        name=key,
+                    )  # type: ignore
+                    for key, _ in expectation.kwargs.items()  # type: ignore
+                    if key not in {"column", "batch_id"}
+                ]
+
         if "observed_value" not in result["result"]:
             return [
                 TestCaseParameterDefinition(
@@ -361,19 +392,118 @@ class OpenMetadataValidationAction(ValidationAction):
         Returns:
             TestCaseResult: a test case result object
         """
-        try:
-            test_result_value = TestResultValue(
-                name="observed_value",
-                value=str(result["result"]["observed_value"]),
-            )
-        except KeyError:
-            unexpected_percent_total = result["result"].get("unexpected_percent_total")
-            test_result_value = TestResultValue(
-                name="unexpected_percentage_total",
-                value=str(unexpected_percent_total),
+        test_result_values = []
+        result_data = result.get("result", {})
+
+        numeric_fields = [
+            "unexpected_percent",
+            "unexpected_percent_total",
+            "missing_percent",
+            "unexpected_count",
+            "missing_count",
+            "element_count",
+            "observed_value",
+            "success_rate",
+        ]
+
+        for field in numeric_fields:
+            if field in result_data:
+                value = result_data[field]
+
+                if isinstance(value, (int, float)):
+                    test_result_values.append(
+                        TestResultValue(
+                            name=field,
+                            value=str(value),
+                            predictedValue=None,
+                        )
+                    )
+                elif field == "observed_value":
+                    test_result_values.extend(
+                        self._extract_complex_value_from_observed_value(value)
+                    )
+
+        return test_result_values
+
+    def _extract_complex_value_from_observed_value(
+        self, observed_value
+    ) -> List[TestResultValue]:
+        """Extract complex value from observed value
+
+        Args:
+            observed_value: observed value
+        Returns:
+            str: complex value
+        """
+        if isinstance(observed_value, list):
+            return [
+                TestResultValue(
+                    name="element_count",
+                    value=str(len(observed_value)),
+                    predictedValue=None,
+                )
+            ]
+
+        if isinstance(observed_value, dict):
+            if "quantiles" in observed_value:
+                result_values = []
+                quantiles = observed_value["quantiles"]
+                values = observed_value["values"]
+                for quantile, value in zip(quantiles, values):
+                    result_values.append(
+                        TestResultValue(
+                            name=f"quantile_{str(quantile)}",
+                            value=str(value),
+                            predictedValue=None,
+                        )
+                    )
+                return result_values
+
+            # catch all other cases that are not a quantile
+            for k, v in observed_value.items():
+                if isinstance(v, (int, float)):
+                    return [
+                        TestResultValue(
+                            name=k,
+                            value=str(v),
+                            predictedValue=None,
+                        )
+                    ]
+
+        if isinstance(observed_value, str):
+            return [
+                TestResultValue(
+                    name="observed_value",
+                    value=str(1),
+                    predictedValue=None,
+                )
+            ]
+
+        return []
+
+    def _get_expectation_config(
+        self, result: dict
+    ) -> Optional[ExpectationConfiguration]:
+        """Get expectation config from GE test result
+
+        Args:
+            result: GE test result
+        Returns:
+            Optional[ExpectationConfiguration]: expectation configuration
+        """
+        if self.expectation_suite:
+            name = result["expectation_config"]["expectation_type"]
+            expectation = next(
+                (
+                    expectation
+                    for expectation in self.expectation_suite.expectations
+                    if expectation.expectation_type == name
+                ),
+                None,
             )
 
-        return [test_result_value]
+            return expectation
+        return None
 
     def _handle_test_case(self, result: Dict, table_entity: Table):
         """Handle adding test to table entity based on the test case.

metadata/great_expectations/action1xx.py

@@ -319,19 +319,94 @@ class OpenMetadataValidationAction1xx(ValidationAction):
         Returns:
             TestCaseResult: a test case result object
         """
-        try:
-            test_result_value = TestResultValue(
-                name="observed_value",
-                value=str(result["result"]["observed_value"]),
-            )
-        except KeyError:
-            unexpected_percent_total = result["result"].get("unexpected_percent_total")
-            test_result_value = TestResultValue(
-                name="unexpected_percentage_total",
-                value=str(unexpected_percent_total),
-            )
+        test_result_values = []
+        result_data = result.get("result", {})
+
+        numeric_fields = [
+            "unexpected_percent",
+            "unexpected_percent_total",
+            "missing_percent",
+            "unexpected_count",
+            "missing_count",
+            "element_count",
+            "observed_value",
+            "success_rate",
+        ]
+
+        for field in numeric_fields:
+            if field in result_data:
+                value = result_data[field]
+
+                if isinstance(value, (int, float)):
+                    test_result_values.append(
+                        TestResultValue(
+                            name=field,
+                            value=str(value),
+                            predictedValue=None,
+                        )
+                    )
+                elif field == "observed_value":
+                    test_result_values.extend(
+                        self._extract_complex_value_from_observed_value(value)
+                    )
+
+        return test_result_values
+
+    def _extract_complex_value_from_observed_value(
+        self, observed_value
+    ) -> List[TestResultValue]:
+        """Extract complex value from observed value
+
+        Args:
+            observed_value: observed value
+        Returns:
+            str: complex value
+        """
+        if isinstance(observed_value, list):
+            return [
+                TestResultValue(
+                    name="element_count",
+                    value=str(len(observed_value)),
+                    predictedValue=None,
+                )
+            ]
+
+        if isinstance(observed_value, dict):
+            if "quantiles" in observed_value:
+                result_values = []
+                quantiles = observed_value["quantiles"]
+                values = observed_value["values"]
+                for quantile, value in zip(quantiles, values):
+                    result_values.append(
+                        TestResultValue(
+                            name=f"quantile_{str(quantile)}",
+                            value=str(value),
+                            predictedValue=None,
+                        )
+                    )
+                return result_values
+
+            # catch all other cases that are not a quantile
+            for k, v in observed_value.items():
+                if isinstance(v, (int, float)):
+                    return [
+                        TestResultValue(
+                            name=k,
+                            value=str(v),
+                            predictedValue=None,
+                        )
+                    ]
+
+        if isinstance(observed_value, str):
+            return [
+                TestResultValue(
+                    name="observed_value",
+                    value=str(1),
+                    predictedValue=None,
+                )
+            ]
 
-        return [test_result_value]
+        return []
 
     def _handle_test_case(self, result: Dict, table_entity: Table):
         """Handle adding test to table entity based on the test case.

metadata/ingestion/lineage/sql_lineage.py

@@ -14,7 +14,7 @@ Helper functions to handle SQL lineage operations
 import itertools
 import traceback
 from collections import defaultdict
-from typing import Any, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import networkx as nx
 from collate_sqllineage.core.models import Column, DataFunction
@@ -894,7 +894,7 @@ def _get_paths_from_subtree(subtree: DiGraph) -> List[List[Any]]:
 
 
 def get_lineage_by_graph(
-    graph: DiGraph,
+    graph: Optional[DiGraph],
     metadata: OpenMetadata,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
@@ -922,3 +922,27 @@ def get_lineage_by_graph(
         subtree = graph.subgraph(component).copy()
         for path in _get_paths_from_subtree(subtree):
             yield from _process_sequence(path, subtree, metadata)
+
+
+def get_lineage_by_procedure_graph(
+    procedure_graph_map: Optional[Dict],
+    metadata: OpenMetadata,
+) -> Iterable[Either[AddLineageRequest]]:
+    """
+    Generate lineage information from a directed graph.
+    """
+    if procedure_graph_map is None:
+        return
+
+    for procedure_and_procedure_graph in procedure_graph_map.values():
+        for either_lineage in get_lineage_by_graph(
+            graph=procedure_and_procedure_graph.graph,
+            metadata=metadata,
+        ):
+            if either_lineage.left is None and either_lineage.right.edge.lineageDetails:
+                either_lineage.right.edge.lineageDetails.pipeline = EntityReference(
+                    id=procedure_and_procedure_graph.procedure.id,
+                    type="storedProcedure",
+                )
+
+            yield either_lineage

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

@@ -48,6 +48,51 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+class DatabricksEngineWrapper:
+    """Wrapper to store engine and schemas to avoid multiple calls"""
+
+    def __init__(self, engine: Engine):
+        self.engine = engine
+        self.inspector = inspect(engine)
+        self.schemas = None
+        self.first_schema = None
+
+    def get_schemas(self):
+        """Get schemas and cache them"""
+        if self.schemas is None:
+            self.schemas = self.inspector.get_schema_names()
+            if self.schemas:
+                # Find the first schema that's not a system schema
+                for schema in self.schemas:
+                    if schema.lower() not in (
+                        "information_schema",
+                        "performance_schema",
+                        "sys",
+                    ):
+                        self.first_schema = schema
+                        break
+                # If no non-system schema found, use the first one
+                if self.first_schema is None and self.schemas:
+                    self.first_schema = self.schemas[0]
+        return self.schemas
+
+    def get_tables(self):
+        """Get tables using the cached first schema"""
+        if self.first_schema is None:
+            self.get_schemas()  # This will set first_schema
+        if self.first_schema:
+            return self.inspector.get_table_names(self.first_schema)
+        return []
+
+    def get_views(self):
+        """Get views using the cached first schema"""
+        if self.first_schema is None:
+            self.get_schemas()  # This will set first_schema
+        if self.first_schema:
+            return self.inspector.get_view_names(self.first_schema)
+        return []
+
+
 def get_connection_url(connection: DatabricksConnection) -> str:
     url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
     return url
@@ -94,12 +139,14 @@ def test_connection(
         except DatabaseError as soe:
             logger.debug(f"Failed to fetch catalogs due to: {soe}")
 
-    inspector = inspect(connection)
+    # Create wrapper to avoid multiple schema calls
+    engine_wrapper = DatabricksEngineWrapper(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,
+        "GetSchemas": engine_wrapper.get_schemas,
+        "GetTables": engine_wrapper.get_tables,
+        "GetViews": engine_wrapper.get_views,
         "GetDatabases": partial(
             test_database_query,
             engine=connection,

metadata/ingestion/source/database/lineage_source.py

@@ -41,6 +41,7 @@ from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Diale
 from metadata.ingestion.lineage.sql_lineage import (
     get_column_fqn,
     get_lineage_by_graph,
+    get_lineage_by_procedure_graph,
     get_lineage_by_query,
 )
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
@@ -49,7 +50,7 @@ from metadata.ingestion.source.database.query_parser_source import QueryParserSo
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.db_utils import get_view_lineage
-from metadata.utils.filters import filter_by_table
+from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -314,9 +315,19 @@ class LineageSource(QueryParserSource, ABC):
     ) -> Iterable[Either[AddLineageRequest]]:
         try:
             for view in views:
-                if filter_by_table(
-                    self.source_config.tableFilterPattern,
-                    view.table_name,
+                if (
+                    filter_by_database(
+                        self.source_config.databaseFilterPattern,
+                        view.db_name,
+                    )
+                    or filter_by_schema(
+                        self.source_config.schemaFilterPattern,
+                        view.schema_name,
+                    )
+                    or filter_by_table(
+                        self.source_config.tableFilterPattern,
+                        view.table_name,
+                    )
                 ):
                     self.status.filter(
                         view.table_name,
@@ -441,10 +452,21 @@ class LineageSource(QueryParserSource, ABC):
         Based on the query logs, prepare the lineage
         and send it to the sink
         """
+        if (
+            self.procedure_graph_map is None
+            and self.source_config.enableTempTableLineage
+        ):
+            # Create a dictionary to store the directed graph for each procedure
+            self.procedure_graph_map = {}
+
         if self.source_config.processViewLineage:
             yield from self.yield_view_lineage() or []
         if self.source_config.processStoredProcedureLineage:
             yield from self.yield_procedure_lineage() or []
+            yield from get_lineage_by_procedure_graph(
+                procedure_graph_map=self.procedure_graph_map,
+                metadata=self.metadata,
+            )
         if self.source_config.processQueryLineage:
             if hasattr(self.service_connection, "supportsLineageExtraction"):
                 yield from self.yield_query_lineage() or []

metadata/ingestion/source/database/query_parser_source.py

@@ -64,6 +64,7 @@ class QueryParserSource(Source, ABC):
         self.source_config = self.config.sourceConfig.config
         self.start, self.end = get_start_and_end(self.source_config.queryLogDuration)
         self.graph = None
+        self.procedure_graph_map = None
 
         self.engine = None
         if get_engine:

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -19,6 +19,7 @@ from collections import defaultdict
 from datetime import datetime
 from typing import Dict, Iterable, List, Optional, Union
 
+from networkx import DiGraph
 from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.engine import Engine
 
@@ -41,7 +42,12 @@ from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.filters import filter_by_stored_procedure
+from metadata.utils.filters import (
+    filter_by_database,
+    filter_by_schema,
+    filter_by_stored_procedure,
+)
+from metadata.utils.helpers import pprint_format_object
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.stored_procedures import get_procedure_name_from_call
 from metadata.utils.time_utils import datetime_to_timestamp
@@ -80,6 +86,17 @@ class ProcedureAndQuery(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
 
+class ProcedureAndProcedureGraph(BaseModel):
+    """
+    Model to hold the procedure and its graph
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    procedure: StoredProcedure
+    graph: DiGraph
+
+
 class StoredProcedureLineageMixin(ABC):
     """
     The full flow is:
@@ -137,6 +154,17 @@ class StoredProcedureLineageMixin(ABC):
                     )
                 )
 
+        queries_count_per_procedure = {
+            procedure_name: len(queries)
+            for procedure_name, queries in queries_dict.items()
+        }
+        logger.info(
+            f"Count of queries executed for stored procedures: {sum(queries_count_per_procedure.values())}"
+        )
+        logger.info(
+            f"Count of queries per stored procedure: {pprint_format_object(queries_count_per_procedure)}"
+        )
+
         return queries_dict
 
     @staticmethod
@@ -166,6 +194,18 @@ class StoredProcedureLineageMixin(ABC):
             query_type=query_by_procedure.query_type,
             query_text=query_by_procedure.query_text,
         ):
+            graph = None
+            if self.source_config.enableTempTableLineage:
+                if not self.procedure_graph_map.get(procedure.fullyQualifiedName.root):
+                    # Map to store the directed graph for each procedure with its FQN as key
+                    self.procedure_graph_map[
+                        procedure.fullyQualifiedName.root
+                    ] = ProcedureAndProcedureGraph(procedure=procedure, graph=DiGraph())
+
+                graph = self.procedure_graph_map.get(
+                    procedure.fullyQualifiedName.root
+                ).graph
+
             self.stored_procedure_query_lineage = True
             for either_lineage in get_lineage_by_query(
                 self.metadata,
@@ -178,6 +218,7 @@ class StoredProcedureLineageMixin(ABC):
                 ),
                 timeout_seconds=self.source_config.parsingTimeoutLimit,
                 lineage_source=LineageSource.QueryLineage,
+                graph=graph,
             ):
                 if (
                     either_lineage.left is None
@@ -290,9 +331,19 @@ class StoredProcedureLineageMixin(ABC):
             or []
         ):
             if procedure:
-                if filter_by_stored_procedure(
-                    self.source_config.storedProcedureFilterPattern,
-                    procedure.name.root,
+                if (
+                    filter_by_database(
+                        self.source_config.databaseFilterPattern,
+                        procedure.database.name,
+                    )
+                    or filter_by_schema(
+                        self.source_config.schemaFilterPattern,
+                        procedure.databaseSchema.name,
+                    )
+                    or filter_by_stored_procedure(
+                        self.source_config.storedProcedureFilterPattern,
+                        procedure.name.root,
+                    )
                 ):
                     self.status.filter(
                         procedure.name.root,

metadata/utils/helpers.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import hashlib
 import itertools
+import pprint
 import re
 import shutil
 import sys
@@ -568,3 +569,10 @@ def evaluate_threshold(threshold: int, operator: str, result: int) -> bool:
         f"Invalid threshold: {threshold}, "
         "Allowed format: <, >, <=, >=, ==, !=. Example: >5"
     )
+
+
+def pprint_format_object(data: Any) -> str:
+    """
+    Pretty print an object in a format that is easy to read
+    """
+    return pprint.pformat(data, width=150)