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

metadata/great_expectations/action1xx.py

@@ -0,0 +1,395 @@
+#  Copyright 2022 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Great Expectations subpackage to send expectation results to
+Open Metadata table quality.
+
+This subpackage needs to be used in Great Expectations
+checkpoints actions.
+"""
+import logging
+import traceback
+from datetime import datetime
+from typing import Dict, List, Literal, Optional, Union, cast
+
+from great_expectations.checkpoint import (
+    ActionContext,
+    CheckpointResult,
+    ValidationAction,
+)
+from great_expectations.core.batch import Batch
+from great_expectations.core.expectation_validation_result import (
+    ExpectationSuiteValidationResultMeta,
+)
+from great_expectations.datasource.fluent import DataAsset
+from great_expectations.validator.validator import Validator
+from sqlalchemy.engine.base import Connection, Engine
+from sqlalchemy.engine.url import URL
+
+from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteRequest
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.tests.basic import (
+    TestCaseResult,
+    TestCaseStatus,
+    TestResultValue,
+)
+from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
+from metadata.generated.schema.tests.testDefinition import (
+    EntityType,
+    TestCaseParameterDefinition,
+    TestPlatform,
+)
+from metadata.generated.schema.tests.testSuite import TestSuite
+from metadata.generated.schema.type.basic import Timestamp
+from metadata.great_expectations.utils.ometa_config_handler import (
+    create_jinja_environment,
+    create_ometa_connection_obj,
+    render_template,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils import fqn
+from metadata.utils.entity_link import get_entity_link
+
+logger = logging.getLogger(
+    "great_expectations.validation_operators.validation_operators.openmetadata"
+)
+
+
+class OpenMetadataValidationAction1xx(ValidationAction):
+    """Open Metdata validation action for GE 1.x.x It inherits from
+    great expection validation action class and implements the
+    `run` method.
+
+    Attributes:
+        data_context: great expectation data context
+        database_service_name: name of the service for the table
+        api_version: default to v1
+        config_file_path: path to the open metdata config path
+    """
+
+    type: Literal["open_metadata_validation_action"] = "open_metadata_validation_action"
+    name: str = "OpenMetadataValidationAction"
+    config_file_path: Optional[str] = None
+    database_service_name: Optional[str] = None
+    schema_name: Optional[str] = "default"
+    database_name: str
+    table_name: Optional[str] = None
+    # Using Optional to make this field not part of the serialized model
+    # This will be initialized in the run method
+    ometa_conn: Optional[OpenMetadata] = None
+
+    def run(  # pylint: disable=unused-argument, arguments-differ
+        self,
+        checkpoint_result: CheckpointResult,
+        action_context: Union[ActionContext, None],
+    ):
+        """main function to implement great expectation hook
+
+        Args:
+            validation_result_suite: result suite returned when checkpoint is ran
+            validation_result_suite_identifier: type of result suite
+            data_asset:
+            expectation_suite_identifier: type of expectation suite
+            checkpoint_identifier: identifier for the checkpoint
+        """
+        self.ometa_conn = self._create_ometa_connection()
+        for _, v in checkpoint_result.run_results.items():
+            meta = v.meta
+            if meta:
+                check_point_spec = self._get_checkpoint_batch_spec(meta)
+                table_entity = self._get_table_entity(
+                    self.database_name,
+                    check_point_spec.get("schema_name", self.schema_name),
+                    check_point_spec.get("table_name"),
+                )
+
+            else:
+                table_entity = self._get_table_entity(
+                    self.database_name,
+                    self.schema_name,
+                    self.table_name,
+                )
+
+            if table_entity:
+                for result in v.results:
+                    self._handle_test_case(result, table_entity)
+
+    @staticmethod
+    def _get_checkpoint_batch_spec(
+        meta: Union[ExpectationSuiteValidationResultMeta, dict]
+    ):
+        """Return run meta and check instance of data_asset
+
+        Args:
+            data_asset: data assets of the checkpoint run
+        Returns:
+            SqlAlchemyDatasourceBatchSpec
+        Raises:
+            ValueError: if datasource not SqlAlchemyDatasourceBatchSpec raise
+        """
+        return meta.get("batch_spec")
+
+    def _get_table_entity(
+        self,
+        database: Optional[str],
+        schema_name: Optional[str],
+        table_name: Optional[str],
+    ) -> Optional[Table]:
+        """Return the table entity for the test. If service name is defined
+        in GE checkpoint entity will be fetch using the FQN. If not provided
+        iterative search will be perform among all the entities. If 2 entities
+        are found with the same `database`.`schema`.`table` the method will
+        raise an error.
+
+        Args:
+            database: database name
+            schema_name: schema name
+            table_name: table name
+
+        Return:
+           Optional[Table]
+
+        Raises:
+             ValueError: if 2 entities with the same
+                         `database`.`schema`.`table` are found
+        """
+        if not all([schema_name, table_name]):
+            raise ValueError(
+                "No Schema or Table name provided. Can't fetch table entity from OpenMetadata."
+            )
+
+        if self.database_service_name:
+            return self.ometa_conn.get_by_name(
+                entity=Table,
+                fqn=f"{self.database_service_name}.{database}.{schema_name}.{table_name}",
+                fields=["testSuite"],
+            )
+
+        table_entity = [
+            entity
+            for entity in self.ometa_conn.list_entities(
+                entity=Table, fields=["testSuite"]
+            ).entities
+            if f"{database}.{schema_name}.{table_name}"
+            in entity.fullyQualifiedName.root
+        ]
+
+        if len(table_entity) > 1:
+            raise ValueError(
+                f"Non unique `database`.`schema`.`table` found: {table_entity}."
+                "Please specify an `database_service_name` in you checkpoint.yml file.",
+            )
+
+        if table_entity:
+            return table_entity[0]
+
+        logger.warning(
+            "No entity found for %s.%s.%s", database, schema_name, table_name
+        )
+        return None
+
+    def _check_or_create_test_suite(self, table_entity: Table) -> TestSuite:
+        """Check if test suite already exists for a given table entity. If not
+        create a new one.
+
+        Args:
+            table_entity: table entity object
+        Returns:
+            TestSuite
+        """
+
+        if table_entity.testSuite:
+            test_suite = self.ometa_conn.get_by_name(
+                TestSuite, table_entity.testSuite.fullyQualifiedName
+            )
+            test_suite = cast(TestSuite, test_suite)
+            return test_suite
+
+        create_test_suite = CreateTestSuiteRequest(
+            name=f"{table_entity.fullyQualifiedName.root}.TestSuite",
+            basicEntityReference=table_entity.fullyQualifiedName.root,
+        )  # type: ignore
+        test_suite = self.ometa_conn.create_or_update_executable_test_suite(
+            create_test_suite
+        )
+        return test_suite
+
+    @staticmethod
+    def _get_execution_engine_url(
+        data_asset: Union[Validator, DataAsset, Batch]
+    ) -> URL:
+        """Get execution engine used to run the expectation
+
+        Args:
+            data_asset: data assets of the checkpoint run
+        Returns:
+            URL
+        Raises:
+            ValueError: if expectation is not ran against DB
+        """
+        if isinstance(data_asset.execution_engine.engine, Engine):
+            return data_asset.execution_engine.engine.url
+        if isinstance(data_asset.execution_engine.engine, Connection):
+            return data_asset.execution_engine.engine.engine.url
+        raise ValueError(
+            "Type is not supported. Make sur you ran your"
+            " expectations against a relational database"
+        )
+
+    def _create_ometa_connection(self) -> OpenMetadata:
+        """Create OpenMetadata API connection"""
+        environment = create_jinja_environment(self.config_file_path)
+        rendered_config = render_template(environment)
+
+        return OpenMetadata(create_ometa_connection_obj(rendered_config))
+
+    def _build_test_case_fqn(self, table_fqn: str, result: Dict) -> str:
+        """build test case fqn from table entity and GE test results
+
+        Args:
+            table_fqn (str): table fully qualified name
+            result (Dict): result from great expectation tests
+        """
+        split_table_fqn = table_fqn.split(".")
+        fqn_ = fqn.build(
+            self.ometa_conn,
+            entity_type=TestCase,
+            service_name=split_table_fqn[0],
+            database_name=split_table_fqn[1],
+            schema_name=split_table_fqn[2],
+            table_name=split_table_fqn[3],
+            column_name=result["expectation_config"]["kwargs"].get("column"),
+            test_case_name=result["expectation_config"]["type"],
+        )
+        fqn_ = cast(str, fqn_)
+        return fqn_
+
+    def _get_test_case_params_value(self, result: dict) -> List[TestCaseParameterValue]:
+        """Build test case parameter value from GE test result"""
+        if "observed_value" not in result["result"]:
+            return [
+                TestCaseParameterValue(
+                    name="unexpected_percentage_total",
+                    value=str(0.0),
+                )
+            ]
+
+        return [
+            TestCaseParameterValue(
+                name=key,
+                value=str(value),
+            )
+            for key, value in result["expectation_config"]["kwargs"].items()
+            if key not in {"column", "batch_id"}
+        ]
+
+    def _get_test_case_params_definition(
+        self, result: dict
+    ) -> List[TestCaseParameterDefinition]:
+        """Build test case parameter definition from GE test result"""
+        if "observed_value" not in result["result"]:
+            return [
+                TestCaseParameterDefinition(
+                    name="unexpected_percentage_total",
+                )  # type: ignore
+            ]
+
+        return [
+            TestCaseParameterDefinition(
+                name=key,
+            )  # type: ignore
+            for key, _ in result["expectation_config"]["kwargs"].items()
+            if key not in {"column", "batch_id"}
+        ]
+
+    def _get_test_result_value(self, result: dict) -> List[TestResultValue]:
+        """Get test result value from GE test result
+
+        Args:
+            result (dict): result
+
+        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),
+            )
+
+        return [test_result_value]
+
+    def _handle_test_case(self, result: Dict, table_entity: Table):
+        """Handle adding test to table entity based on the test case.
+        Test Definitions will be created on the fly from the results of the
+        great expectations run. We will then write the test case results to the
+        specific test case.
+
+        Args:
+            result: GE test result
+            table_entity: table entity object
+        """
+
+        try:
+            test_definition = self.ometa_conn.get_or_create_test_definition(
+                test_definition_fqn=result["expectation_config"]["type"],
+                test_definition_description=result["expectation_config"][
+                    "type"
+                ].replace("_", " "),
+                entity_type=EntityType.COLUMN
+                if "column" in result["expectation_config"]["kwargs"]
+                else EntityType.TABLE,
+                test_platforms=[TestPlatform.GreatExpectations],
+                test_case_parameter_definition=self._get_test_case_params_definition(
+                    result
+                ),
+            )
+
+            test_case_fqn = self._build_test_case_fqn(
+                table_entity.fullyQualifiedName.root,
+                result,
+            )
+
+            test_case = self.ometa_conn.get_or_create_test_case(
+                test_case_fqn,
+                entity_link=get_entity_link(
+                    Table,
+                    fqn=table_entity.fullyQualifiedName.root,
+                    column_name=fqn.split_test_case_fqn(test_case_fqn).column,
+                ),
+                test_definition_fqn=test_definition.fullyQualifiedName.root,
+                test_case_parameter_values=self._get_test_case_params_value(result),
+            )
+
+            self.ometa_conn.add_test_case_results(
+                test_results=TestCaseResult(
+                    timestamp=Timestamp(int(datetime.now().timestamp() * 1000)),
+                    testCaseStatus=TestCaseStatus.Success
+                    if result["success"]
+                    else TestCaseStatus.Failed,
+                    testResultValue=self._get_test_result_value(result),
+                ),  # type: ignore
+                test_case_fqn=test_case.fullyQualifiedName.root,
+            )
+
+            logger.debug(
+                f"Test case result for {test_case.fullyQualifiedName.root} successfully ingested"
+            )
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(exc)

metadata/ingestion/connections/builders.py

@@ -31,6 +31,7 @@ from metadata.generated.schema.entity.services.connections.database.common.iamAu
     IamAuthConfigurationSource,
 )
 from metadata.ingestion.connections.headers import inject_query_header_by_conn
+from metadata.ingestion.connections.query_logger import attach_query_tracker
 from metadata.ingestion.connections.secrets import connection_with_options_secrets
 from metadata.utils.constants import BUILDER_PASSWORD_ATTR
 from metadata.utils.logger import cli_logger
@@ -76,6 +77,8 @@ def create_generic_db_connection(
         max_overflow=-1,
     )
 
+    attach_query_tracker(engine)
+
     if hasattr(connection, "supportsQueryComment"):
         listen(
             engine,
@@ -125,11 +128,10 @@ def init_empty_connection_options() -> ConnectionOptions:
     return ConnectionOptions(root={})
 
 
-def _add_password(url: str, connection) -> str:
+def get_password_secret(connection) -> SecretStr:
     """
-    A helper function that adds the password to the url if it exists.
-    Distinguishing between BasicAuth (Password) and IamAuth (AWSConfig)
-    and adding to url.
+    Helper to extract the password secret from the connection object.
+    Handles BasicAuth, IamAuth, and falls back to empty SecretStr if not found.
     """
     password = getattr(connection, BUILDER_PASSWORD_ATTR, None)
 
@@ -158,6 +160,14 @@ def _add_password(url: str, connection) -> str:
     if not password:
         logger.warning("No password has been provided in connection")
         password = SecretStr("")
+    return password
+
+
+def _add_password(url: str, connection) -> str:
+    """
+    A helper function that adds the password to the url if it exists.
+    """
+    password = get_password_secret(connection)
     url += f":{quote_plus(password.get_secret_value())}"
     return url
 

metadata/ingestion/connections/connection.py

@@ -0,0 +1,63 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+BaseConnection abstract class for database connectors.
+
+This module defines the BaseConnection abstract base class,
+which provides a common, type-safe interface for all database connection implementations.
+Each connector subclass should inherit from BaseConnection and implement the required abstract methods
+to provide a unified way to instantiate and interact with different data sources.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Generic, Optional, TypeVar
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
+
+S = TypeVar("S")  # ServiceConnection Type
+C = TypeVar("C")  # Client Type
+
+
+class BaseConnection(ABC, Generic[S, C]):
+    """
+    Abstract base class for database connections, providing a unified interface
+    for service connection and client/engine access across different data sources.
+    """
+
+    service_connection: S
+
+    def __init__(self, service_connection: S) -> None:
+        self.service_connection = service_connection
+
+    @abstractmethod
+    def get_client(self) -> C:
+        """
+        Return the main client/engine/connection object for this service.
+        """
+
+    @abstractmethod
+    def test_connection(
+        self,
+        metadata: OpenMetadata,
+        automation_workflow: Optional[AutomationWorkflow] = None,
+        timeout_seconds: Optional[int] = THREE_MIN,
+    ) -> TestConnectionResult:
+        """
+        Test the connection to the service.
+        """

metadata/ingestion/connections/query_logger.py

@@ -0,0 +1,115 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  pylint: disable=W0613
+
+"""
+Query tracking implementation using SQLAlchemy event listeners
+"""
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional, Tuple, Union
+
+from pydantic import BaseModel, ConfigDict
+from sqlalchemy.event import listen
+from sqlalchemy.sql.elements import TextClause
+
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class QueryInfo(BaseModel):
+    """Class to store information about a query execution"""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    statement: Union[str, TextClause]
+    parameters: Optional[Union[Dict[str, Any], Tuple[Any, ...]]]
+    start_time: datetime
+    end_time: Optional[datetime] = None
+    duration_ms: Optional[float] = None
+    error: Optional[Exception] = None
+
+
+class QueryLogger:
+    """Class to track SQL query execution using SQLAlchemy event listeners"""
+
+    def __init__(self):
+        self._current_query: Optional[QueryInfo] = None
+
+    def before_cursor_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: Union[str, TextClause],
+        parameters: Optional[Dict[str, Any]],
+        context: Any,
+        executemany: bool,
+    ) -> Tuple[Union[str, TextClause], Optional[Dict[str, Any]]]:
+        """Event listener for before cursor execute"""
+        self._current_query = QueryInfo(
+            statement=statement,
+            parameters=parameters,
+            start_time=datetime.now(timezone.utc),
+        )
+        return statement, parameters
+
+    def after_cursor_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: Union[str, TextClause],
+        parameters: Optional[Dict[str, Any]],
+        context: Any,
+        executemany: bool,
+    ) -> None:
+        """Event listener for after cursor execute"""
+        if self._current_query:
+            query = self._current_query
+            query.end_time = datetime.now(timezone.utc)
+            query.duration_ms = (
+                query.end_time - query.start_time
+            ).total_seconds() * 1000
+
+            logger.debug(
+                "Query execution details:\n"
+                f"  Start Time: {query.start_time}\n"
+                f"  End Time: {query.end_time}\n"
+                f"  Duration: {query.duration_ms:.2f} ms\n"
+                f"  Query: {query.statement}\n"
+                f"  Parameters: {query.parameters}"
+            )
+
+            self._current_query = None
+
+
+def attach_query_tracker(engine: Any):
+    """
+    Attach query tracking event listeners to a SQLAlchemy engine
+
+    Args:
+        engine: SQLAlchemy engine to attach listeners to
+
+    Returns:
+        QueryLogger instance that can be used to access query execution data
+    """
+    tracker = QueryLogger()
+
+    listen(
+        engine,
+        "before_cursor_execute",
+        tracker.before_cursor_execute,
+        retval=True,
+    )
+    listen(
+        engine,
+        "after_cursor_execute",
+        tracker.after_cursor_execute,
+    )

metadata/ingestion/connections/test_connections.py

@@ -38,7 +38,7 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
 )
 from metadata.generated.schema.type.basic import Timestamp
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import kill_active_connections
+from metadata.ingestion.source.connections_utils import kill_active_connections
 from metadata.profiler.orm.functions.conn_test import ConnTestFn
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import cli_logger

metadata/ingestion/ometa/mixins/tests_mixin.py

@@ -74,8 +74,8 @@ class OMetaTestsMixin:
         Returns:
             _type_: _description_
         """
-        resp = self.client.put(
-            f"{self.get_suffix(TestCase)}/{quote(test_case_fqn)}/testCaseResult",
+        resp = self.client.post(
+            f"{self.get_suffix(TestCaseResult)}/{quote(test_case_fqn)}",
             test_results.model_dump_json(),
         )
 
@@ -167,7 +167,6 @@ class OMetaTestsMixin:
         self,
         test_case_fqn: str,
         entity_link: Optional[str] = None,
-        test_suite_fqn: Optional[str] = None,
         test_definition_fqn: Optional[str] = None,
         test_case_parameter_values: Optional[List[TestCaseParameterValue]] = None,
     ):
@@ -196,7 +195,6 @@ class OMetaTestsMixin:
             CreateTestCaseRequest(
                 name=test_case_fqn.split(".")[-1],
                 entityLink=entity_link,
-                testSuite=test_suite_fqn,
                 testDefinition=test_definition_fqn,
                 parameterValues=test_case_parameter_values,
             )  # type: ignore
@@ -253,7 +251,7 @@ class OMetaTestsMixin:
         }
 
         resp = self.client.get(
-            f"/dataQuality/testCases/{test_case_fqn}/testCaseResult",
+            f"{self.get_suffix(TestCaseResult)}/{test_case_fqn}",
             params,
         )
 

metadata/ingestion/ometa/routes.py

@@ -167,6 +167,7 @@ from metadata.generated.schema.entity.teams.role import Role
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import AuthenticationMechanism, User
 from metadata.generated.schema.settings.settings import Settings
+from metadata.generated.schema.tests.basic import TestCaseResult
 from metadata.generated.schema.tests.testCase import TestCase
 from metadata.generated.schema.tests.testDefinition import TestDefinition
 from metadata.generated.schema.tests.testSuite import TestSuite
@@ -264,6 +265,7 @@ ROUTES = {
     TestSuite.__name__: "/dataQuality/testSuites",
     CreateTestSuiteRequest.__name__: "/dataQuality/testSuites",
     TestCase.__name__: "/dataQuality/testCases",
+    TestCaseResult.__name__: "/dataQuality/testCases/testCaseResults",
     CreateTestCaseRequest.__name__: "/dataQuality/testCases",
     # Analytics
     WebAnalyticEventData.__name__: "/analytics/web/events/collect",