openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/ingestion/ometa/mixins/logs_mixin.py

@@ -0,0 +1,316 @@
+#  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.
+"""
+Mixin class for centralized log streaming to S3 storage
+
+This mixin provides methods to send ingestion logs to OpenMetadata server's
+S3 storage backend. It's designed to be used by OpenMetadata class and
+integrates with the existing logging infrastructure.
+"""
+
+import base64
+import gzip
+import json
+import os
+import socket
+import time
+from typing import Optional
+from uuid import UUID
+
+from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.utils import model_str
+from metadata.utils.constants import UTF_8
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaLogsMixin:
+    """
+    OpenMetadata API methods for log streaming to S3.
+
+    This mixin provides centralized methods for shipping logs to the
+    OpenMetadata server's S3 storage backend. It reuses the existing
+    OpenMetadata client instance to avoid creating new sessions.
+    """
+
+    client: REST
+
+    def send_logs_to_s3(
+        self,
+        pipeline_fqn: str,
+        run_id: UUID,
+        log_content: str,
+        compress: bool = False,
+    ) -> bool:
+        """
+        Send logs to S3 storage via OpenMetadata server endpoint.
+
+        Args:
+            pipeline_fqn: Fully qualified name of the ingestion pipeline
+            run_id: Unique identifier for the pipeline run
+            log_content: The log content to send
+            compress: Whether to compress logs before sending
+
+        Returns:
+            bool: True if logs were sent successfully, False otherwise
+        """
+        try:
+            # Extract the UUID string value from the object if it has a .root attribute
+            # Build the API endpoint
+            url = (
+                f"/services/ingestionPipelines/logs/{pipeline_fqn}/{model_str(run_id)}"
+            )
+
+            # Prepare log batch data matching Java LogBatch structure
+            log_batch = {
+                "logs": log_content,
+                "timestamp": int(time.time() * 1000),  # milliseconds since epoch
+                "connectorId": f"{socket.gethostname()}-{os.getpid()}",
+                "compressed": False,
+                "lineCount": log_content.count("\n") + 1,
+            }
+
+            # Apply compression if requested and content is large enough
+            if compress and len(log_content) > 10240:  # Compress if > 10KB
+                compressed_data = gzip.compress(log_content.encode(UTF_8))
+                log_batch["logs"] = base64.b64encode(compressed_data).decode(UTF_8)
+                log_batch["compressed"] = True
+
+            # Send logs using the existing client
+            self.client.post(
+                url,
+                data=json.dumps(log_batch),
+            )
+
+            # The REST client returns None for successful requests with empty response body (HTTP 200/201/204)
+            # If we reach this point without an exception, the request was successful
+            logger.debug(
+                f"Successfully sent {log_batch['lineCount']} log lines for pipeline {pipeline_fqn}"
+            )
+            return True
+
+        except Exception as e:
+            line_count = log_content.count("\n") + 1
+            logger.error(
+                f"Failed to send logs to S3 for pipeline {pipeline_fqn}: {e}. "
+                f"Log batch size: {len(log_content)} chars, "
+                f"Lines: {line_count}, "
+                f"Compressed: {log_batch.get('compressed', False)}"
+            )
+            return False
+
+    def send_logs_batch(
+        self,
+        pipeline_fqn: str,
+        run_id: UUID,
+        log_content: str,
+        enable_compression: bool = False,
+        max_retries: int = 3,
+    ) -> dict:
+        """
+        Send logs batch to S3 storage via OpenMetadata server endpoint with retry logic.
+
+        Args:
+            pipeline_fqn: Fully qualified name of the ingestion pipeline
+            run_id: Unique identifier for the pipeline run
+            log_content: The log content to send
+            enable_compression: Whether to compress logs before sending
+            max_retries: Maximum number of retry attempts
+
+        Returns:
+            dict: Metrics including lines sent and bytes sent
+        """
+        metrics = {"logs_sent": 0, "bytes_sent": 0}
+
+        for attempt in range(max_retries + 1):
+            try:
+                success = self.send_logs_to_s3(
+                    pipeline_fqn=pipeline_fqn,
+                    run_id=run_id,
+                    log_content=log_content,
+                    compress=enable_compression and len(log_content) > 10240,
+                )
+
+                if success:
+                    # Update metrics
+                    line_count = log_content.count("\n") + 1
+                    metrics["logs_sent"] = line_count
+                    metrics["bytes_sent"] = len(log_content)
+
+                    if attempt > 0:
+                        logger.info(
+                            f"Successfully shipped {line_count} log lines to server on attempt {attempt + 1}"
+                        )
+                    else:
+                        logger.debug(
+                            f"Successfully shipped {line_count} log lines to server"
+                        )
+                    return metrics
+                else:
+                    if attempt < max_retries:
+                        wait_time = 2**attempt  # Exponential backoff: 1s, 2s, 4s
+                        logger.warning(
+                            f"Failed to send logs for pipeline {pipeline_fqn}, "
+                            f"attempt {attempt + 1}/{max_retries + 1}. Retrying in {wait_time}s..."
+                        )
+                        time.sleep(wait_time)
+                    else:
+                        logger.error(
+                            f"Failed to send logs for pipeline {pipeline_fqn} after {max_retries + 1} attempts"
+                        )
+
+            except Exception as e:
+                if attempt < max_retries:
+                    wait_time = 2**attempt  # Exponential backoff: 1s, 2s, 4s
+                    logger.warning(
+                        f"Error sending logs batch for pipeline {pipeline_fqn}: {e}. "
+                        f"Attempt {attempt + 1}/{max_retries + 1}. Retrying in {wait_time}s..."
+                    )
+                    time.sleep(wait_time)
+                else:
+                    logger.error(
+                        f"Failed to send logs batch for pipeline {pipeline_fqn} after {max_retries + 1} attempts: {e}"
+                    )
+
+        return metrics
+
+    def create_log_stream(
+        self,
+        pipeline_fqn: str,
+        run_id: UUID,
+    ) -> None:
+        """
+        Initialize a log stream for a pipeline run.
+
+        This method can be used to set up a log stream session with the server.
+
+        Args:
+            pipeline_fqn: Fully qualified name of the ingestion pipeline
+            run_id: Unique identifier for the pipeline run
+        """
+        try:
+
+            # Initialize log stream with the server
+            url = (
+                f"/services/ingestionPipelines/logs/{pipeline_fqn}/{model_str(run_id)}"
+            )
+
+            init_data = {
+                "connectorId": f"{socket.gethostname()}-{os.getpid()}",
+                "timestamp": int(time.time() * 1000),
+            }
+
+            response = self.client.post(
+                url,
+                data=json.dumps(init_data),
+            )
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to initialize log stream for pipeline {pipeline_fqn}: {e}"
+            )
+
+    def close_log_stream(
+        self,
+        pipeline_fqn: str,
+        run_id: UUID,
+    ) -> bool:
+        """
+        Close a log stream for a pipeline run.
+
+        This method signals the server that log streaming has completed
+        for a particular pipeline run.
+
+        Args:
+            pipeline_fqn: Fully qualified name of the ingestion pipeline
+            run_id: Unique identifier for the pipeline run
+
+        Returns:
+            bool: True if stream was closed successfully, False otherwise
+        """
+        try:
+
+            url = f"/services/ingestionPipelines/logs/{pipeline_fqn}/{model_str(run_id)}/close"
+
+            close_data = {
+                "connectorId": f"{socket.gethostname()}-{os.getpid()}",
+                "timestamp": int(time.time() * 1000),
+            }
+
+            self.client.post(
+                url,
+                data=json.dumps(close_data),
+            )
+
+            logger.debug(f"Successfully closed log stream for pipeline {pipeline_fqn}")
+
+            return True
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to close log stream for pipeline {pipeline_fqn}: {e}"
+            )
+            return False
+
+    def get_logs_from_s3(
+        self,
+        pipeline_fqn: str,
+        run_id: UUID,
+        offset: int = 0,
+        limit: int = 1000,
+    ) -> Optional[str]:
+        """
+        Retrieve logs from S3 storage for a pipeline run.
+
+        Args:
+            pipeline_fqn: Fully qualified name of the ingestion pipeline
+            run_id: Unique identifier for the pipeline run
+            offset: Starting offset for log retrieval
+            limit: Maximum number of log lines to retrieve
+
+        Returns:
+            Optional[str]: Log content if available, None otherwise
+        """
+        try:
+
+            url = (
+                f"/services/ingestionPipelines/logs/{pipeline_fqn}/{model_str(run_id)}"
+            )
+
+            params = {
+                "offset": offset,
+                "limit": limit,
+            }
+
+            response = self.client.get(url, data=params)
+
+            if response and isinstance(response, dict):
+                log_data = response.get("logs", "")
+
+                # Decompress if necessary
+                if response.get("compressed"):
+                    try:
+                        decoded = base64.b64decode(log_data)
+                        log_data = gzip.decompress(decoded).decode(UTF_8)
+                    except Exception as e:
+                        logger.error(f"Failed to decompress logs: {e}")
+                        return None
+
+                return log_data
+
+            return None
+
+        except Exception as e:
+            logger.error(
+                f"Failed to retrieve logs from S3 for pipeline {pipeline_fqn}: {e}"
+            )
+            return None

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -20,6 +20,9 @@ from typing import Dict, List, Optional, Type, TypeVar, Union
 
 from pydantic import BaseModel
 
+from metadata.generated.schema.entity.automations.response.queryRunnerResponse import (
+    QueryRunnerResponse,
+)
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -128,6 +131,11 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         Given an Entity type and Source entity and Destination entity,
         generate a JSON Patch and apply it.
 
+        This method provides fine-grained control over entity updates and can
+        override fields that create_or_update() cannot due to server-side restrictions
+        across various entity types. Use override_metadata=True to force updates of
+        protected metadata fields.
+
         Args
             entity (T): Entity Type
             source: Source payload which is current state of the source in OpenMetadata
@@ -135,7 +143,8 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             allowed_fields: List of field names to filter from source and destination models
             restrict_update_fields: List of field names which will only support add operation
             array_entity_fields: List of array fields to sort for consistent patching
-            override_metadata: Whether to override existing metadata fields
+            override_metadata: Whether to override existing metadata fields. Set to True
+                to force updates of protected fields across various entity types.
             skip_on_failure: Whether to skip the patch operation on failure (default: True)
 
         Returns
@@ -188,12 +197,16 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         """
         Given an Entity type and ID, JSON PATCH the description.
 
+        This method is useful when you need to update descriptions that cannot be
+        overridden through create_or_update() due to server-side business rules
+        across various entity types. Use force=True to override existing descriptions.
+
         Args
             entity (T): Entity Type
             source: source entity object
             description: new description to add
             force: if True, we will patch any existing description. Otherwise, we will maintain
-                the existing data.
+                the existing data. Set to True to override existing descriptions across entity types.
             skip_on_failure: if True, return None on failure instead of raising exception
         Returns
             Updated Entity
@@ -547,7 +560,9 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
     def patch_automation_workflow_response(
         self,
         automation_workflow: AutomationWorkflow,
-        result: Union[TestConnectionResult, ReverseIngestionResponse],
+        result: Union[
+            TestConnectionResult, ReverseIngestionResponse, QueryRunnerResponse
+        ],
         workflow_status: WorkflowStatus,
     ) -> None:
         """
@@ -560,17 +575,16 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         }
 
         # for deserializing into json convert enum object to string
-        if isinstance(result, TestConnectionResult):
-            result_data[PatchField.VALUE]["status"] = result_data[PatchField.VALUE][
-                "status"
-            ].value
-        else:
+        if isinstance(result, ReverseIngestionResponse):
             # Convert UUID in string
             data = result_data[PatchField.VALUE]
             data["serviceId"] = str(data["serviceId"])
             for operation_result in data["results"]:
                 operation_result["id"] = str(operation_result["id"])
-
+        else:
+            result_data[PatchField.VALUE]["status"] = result_data[PatchField.VALUE][
+                "status"
+            ].value
         status_data: Dict = {
             PatchField.PATH: PatchPath.STATUS,
             PatchField.OPERATION: PatchOperation.ADD,

metadata/ingestion/ometa/mixins/profile_mixin.py

@@ -0,0 +1,77 @@
+#  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.
+"""
+Mixin class containing Pipeline specific methods
+
+To be used by OpenMetadata class
+"""
+
+from typing import Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from metadata.generated.schema.type.entityProfile import EntityProfile, ProfileTypeEnum
+from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.models import EntityList
+from metadata.ingestion.ometa.utils import get_entity_type
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class OMetaProfileMixin:
+    """
+    OpenMetadata API methods related to Profiles.
+
+    To be inherited by OpenMetadata
+    """
+
+    client: REST
+
+    def get_profile_data_by_type(
+        self,
+        entity_type: Type[T],
+        start_ts: int,
+        end_ts: int,
+        profile_type: Optional[ProfileTypeEnum] = None,
+    ) -> EntityList[EntityProfile]:
+        """List all profile data for a given entity type. To get all the profile for
+        a specific profile type use the profile_type parameter.
+
+        Args:
+            entity_type (Type[T]): Entity type
+            start_ts (int): Start timestamp
+            end_ts (int): End timestamp
+            profile_type (Optional[ProfileTypeEnum]): Profile type
+
+        Returns:
+            EntityList[EntityProfile]: EntityList list object
+        """
+        entity_type = get_entity_type(entity_type)
+        params = {
+            "startTs": start_ts,
+            "endTs": end_ts,
+        }
+        if profile_type:
+            params["profileType"] = profile_type.value
+        resp = self.client.get(
+            f"{self.get_suffix(EntityProfile)}/{entity_type}",
+            data=params,
+        )
+        paging = resp["paging"]
+        return EntityList[EntityProfile](
+            entities=[EntityProfile(**entity) for entity in resp["data"]],
+            total=paging["total"],
+            after=paging["after"],
+            before=paging["before"],
+        )

metadata/ingestion/ometa/mixins/tag_glossary_mixin.py

@@ -0,0 +1,68 @@
+#  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.
+"""Tag and Glossary Term specific operations"""
+import traceback
+from typing import Dict
+
+from metadata.ingestion.ometa.client import REST
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaTagGlossaryMixin:
+    """Mixin class containing Tag and Glossary Term specific methods"""
+
+    client: REST
+
+    def get_tag_assets(self, fqn: str, limit: int = 10, offset: int = 0) -> Dict:
+        """
+        Get paginated list of assets for a tag
+
+        Args:
+            fqn: Fully qualified name of the tag
+            limit: Maximum number of assets to return (default 10, max 1000)
+            offset: Offset from which to start returning results (default 0)
+
+        Returns:
+            API response as a dictionary containing paginated assets
+        """
+        try:
+            path = f"/tags/name/{fqn}/assets"
+            params = {"limit": limit, "offset": offset}
+            return self.client.get(path, params)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get tag assets due to {exc}")
+            return {}
+
+    def get_glossary_term_assets(
+        self, fqn: str, limit: int = 10, offset: int = 0
+    ) -> Dict:
+        """
+        Get paginated list of assets for a glossary term
+
+        Args:
+            fqn: Fully qualified name of the glossary term
+            limit: Maximum number of assets to return (default 10, max 1000)
+            offset: Offset from which to start returning results (default 0)
+
+        Returns:
+            API response as a dictionary containing paginated assets
+        """
+        try:
+            path = f"/glossaryTerms/name/{fqn}/assets"
+            params = {"limit": limit, "offset": offset}
+            return self.client.get(path, params)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get glossary term assets due to {exc}")
+            return {}

metadata/ingestion/ometa/mixins/tests_mixin.py

@@ -253,7 +253,7 @@ class OMetaTestsMixin:
         }
 
         resp = self.client.get(
-            f"{self.get_suffix(TestCaseResult)}/{test_case_fqn}",
+            f"{self.get_suffix(TestCaseResult)}/{quote(test_case_fqn)}",
             params,
         )
 

metadata/ingestion/ometa/mixins/user_mixin.py

@@ -14,6 +14,7 @@ Mixin class containing User specific methods
 To be used by OpenMetadata class
 """
 import json
+import traceback
 from functools import lru_cache
 from typing import Optional, Type
 
@@ -181,21 +182,6 @@ class OMetaUserMixin:
         """
         Get a User or Team Entity Reference by searching by its name
         """
-        maybe_user = self._search_by_name(
-            entity=User, name=name, from_count=from_count, size=size, fields=fields
-        )
-        if maybe_user:
-            return EntityReferenceList(
-                root=[
-                    EntityReference(
-                        id=maybe_user.id.root,
-                        type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
-                        name=maybe_user.name.root,
-                        displayName=maybe_user.displayName,
-                    )
-                ]
-            )
-
         maybe_team = self._search_by_name(
             entity=Team, name=name, from_count=from_count, size=size, fields=fields
         )
@@ -213,5 +199,60 @@ class OMetaUserMixin:
                     )
                 ]
             )
-
+        maybe_user = self._search_by_name(
+            entity=User, name=name, from_count=from_count, size=size, fields=fields
+        )
+        if maybe_user:
+            return EntityReferenceList(
+                root=[
+                    EntityReference(
+                        id=maybe_user.id.root,
+                        type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
+                        name=maybe_user.name.root,
+                        displayName=maybe_user.displayName,
+                    )
+                ]
+            )
         return None
+
+    def get_user_assets(self, name: str, limit: int = 10, offset: int = 0) -> dict:
+        """
+        Get paginated list of assets for a user
+
+        Args:
+            name: Name of the user
+            limit: Maximum number of assets to return (default 10, max 1000)
+            offset: Offset from which to start returning results (default 0)
+
+        Returns:
+            API response as a dictionary containing paginated assets
+        """
+        try:
+            path = f"/users/name/{name}/assets"
+            params = {"limit": limit, "offset": offset}
+            return self.client.get(path, params)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get user assets due to {exc}")
+            return {}
+
+    def get_team_assets(self, name: str, limit: int = 10, offset: int = 0) -> dict:
+        """
+        Get paginated list of assets for a team
+
+        Args:
+            name: Name of the team
+            limit: Maximum number of assets to return (default 10, max 1000)
+            offset: Offset from which to start returning results (default 0)
+
+        Returns:
+            API response as a dictionary containing paginated assets
+        """
+        try:
+            path = f"/teams/name/{name}/assets"
+            params = {"limit": limit, "offset": offset}
+            return self.client.get(path, params)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get team assets due to {exc}")
+            return {}