openmetadata-ingestion 1.6.7.0__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -11,16 +11,21 @@
 """
 Base class for ingesting database services
 """
+import traceback
 from abc import ABC, abstractmethod
+from datetime import datetime
 from typing import Any, Iterable, List, Optional, Set
 
-from pydantic import Field
+from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
-from metadata.generated.schema.entity.data.pipeline import Pipeline
+from metadata.generated.schema.entity.data.pipeline import Pipeline, PipelineState
 from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.services.pipelineService import (
     PipelineConnection,
     PipelineService,
@@ -31,6 +36,7 @@ from metadata.generated.schema.metadataIngestion.pipelineServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.usageRequest import UsageRequest
 from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
@@ -54,11 +60,21 @@ from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
 from metadata.ingestion.source.pipeline.openlineage.utils import FQNNotFoundException
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_pipeline
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
+class PipelineUsage(BaseModel):
+    """
+    Wrapper to handle type at the sink
+    """
+
+    pipeline: Pipeline
+    usage: UsageRequest
+
+
 class PipelineServiceTopology(ServiceTopology):
     """
     Defines the hierarchy in Pipeline Services.
@@ -114,6 +130,12 @@ class PipelineServiceTopology(ServiceTopology):
                 consumer=["pipeline_service"],
                 nullable=True,
             ),
+            NodeStage(
+                type_=UsageRequest,
+                processor="yield_pipeline_usage",
+                consumer=["pipeline_service"],
+                nullable=True,
+            ),
         ],
     )
 
@@ -133,6 +155,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     context = TopologyContextManager(topology)
     pipeline_source_state: Set = set()
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,
@@ -141,6 +164,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
+        self.today = datetime.now().strftime("%Y-%m-%d")
         self.service_connection = self.config.serviceConnection.root.config
         self.source_config: PipelineServiceMetadataPipeline = (
             self.config.sourceConfig.config
@@ -182,6 +206,100 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     ) -> Iterable[Either[OMetaPipelineStatus]]:
         """Get Pipeline Status"""
 
+    def get_pipeline_state(self, pipeline_details: Any) -> Optional[PipelineState]:
+        """Get Pipeline State"""
+
+    def yield_pipeline_usage(
+        self, pipeline_details: Any
+    ) -> Iterable[Either[PipelineUsage]]:
+        """
+        Yield the usage of the pipeline
+        we will check the usage of the pipeline
+        by checking the tasks that have run today or are running today or ends today
+        we get the count of these tasks and compare it with the usageSummary
+        if the usageSummary is not present or the date is not today
+        we yield the fresh usage
+        """
+        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: Pipeline = self.metadata.get_by_name(
+                entity=Pipeline,
+                fqn=pipeline_fqn,
+                fields=["tasks", "usageSummary"],
+            )
+
+            if pipeline.tasks:
+                current_task_usage = sum(
+                    1
+                    for task in pipeline.tasks
+                    if task.startDate
+                    and task.startDate.startswith(self.today)
+                    or task.endDate
+                    and task.endDate.startswith(self.today)
+                )
+                if not current_task_usage:
+                    logger.debug(
+                        f"No usage to report for {pipeline.fullyQualifiedName.root}"
+                    )
+
+                if not pipeline.usageSummary:
+                    logger.info(
+                        f"Yielding fresh usage for {pipeline.fullyQualifiedName.root}"
+                    )
+                    yield Either(
+                        right=PipelineUsage(
+                            pipeline=pipeline,
+                            usage=UsageRequest(
+                                date=self.today, count=current_task_usage
+                            ),
+                        )
+                    )
+
+                elif (
+                    str(pipeline.usageSummary.date.root) != self.today
+                    or not pipeline.usageSummary.dailyStats.count
+                ):
+                    latest_usage = pipeline.usageSummary.dailyStats.count
+
+                    new_usage = current_task_usage - latest_usage
+                    if new_usage < 0:
+                        raise ValueError(
+                            f"Wrong computation of usage difference. Got new_usage={new_usage}."
+                        )
+
+                    logger.info(
+                        f"Yielding new usage for {pipeline.fullyQualifiedName.root}"
+                    )
+                    yield Either(
+                        right=PipelineUsage(
+                            pipeline=pipeline,
+                            usage=UsageRequest(date=self.today, count=new_usage),
+                        )
+                    )
+
+                else:
+                    logger.debug(
+                        f"Latest usage {pipeline.usageSummary} vs. today {self.today}. Nothing to compute."
+                    )
+                    logger.info(
+                        f"Usage already informed for {pipeline.fullyQualifiedName.root}"
+                    )
+
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name="Pipeline Usage",
+                    error=f"Exception computing pipeline usage for {pipeline_details}: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
     def yield_pipeline_lineage(
         self, pipeline_details: Any
     ) -> Iterable[Either[OMetaLineageRequest]]:
@@ -229,6 +347,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def close(self):
         """Method to implement any required logic after the ingestion process is completed"""
+        self.metadata.compute_percentile(Pipeline, self.today)
 
     def get_services(self) -> Iterable[WorkflowSource]:
         yield self.config

metadata/ingestion/source/search/opensearch/connection.py

@@ -0,0 +1,215 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+
+"""
+Source connection handler for OpenSearch
+"""
+from pathlib import Path
+from typing import Optional
+
+from opensearchpy import OpenSearch, RequestsHttpConnection
+from requests_aws4auth import AWS4Auth
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.common.sslCertPaths import (
+    SslCertificatesByPath,
+)
+from metadata.generated.schema.entity.services.connections.common.sslCertValues import (
+    SslCertificatesByValues,
+)
+from metadata.generated.schema.entity.services.connections.common.sslConfig import (
+    SslConfig,
+)
+from metadata.generated.schema.entity.services.connections.search.elasticSearch.basicAuth import (
+    BasicAuthentication,
+)
+from metadata.generated.schema.entity.services.connections.search.openSearchConnection import (
+    OpenSearchConnection,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.generated.schema.security.ssl.verifySSLConfig import VerifySSL
+from metadata.ingestion.connections.builders import init_empty_connection_arguments
+from metadata.ingestion.connections.test_connections import test_connection_steps
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN, UTF_8
+from metadata.utils.helpers import clean_uri, init_staging_dir
+
+CA_CERT_FILE_NAME = "root.pem"
+CLIENT_CERT_FILE_NAME = "client.pem"
+KEY_CERT_FILE_NAME = "client_key.pem"
+
+
+def _clean_cert_value(cert_data: str) -> str:
+    return cert_data.replace("\\n", "\n")
+
+
+def write_data_to_file(file_path: Path, cert_data: str) -> None:
+    with open(
+        file_path,
+        "w+",
+        encoding=UTF_8,
+    ) as file:
+        data = _clean_cert_value(cert_data)
+        file.write(data)
+
+
+def _handle_ssl_context_by_value(ssl_config: SslConfig):
+    ca_cert = False
+    client_cert = None
+    private_key = None
+    init_staging_dir(ssl_config.certificates.stagingDir)
+    if ssl_config.certificates.caCertValue:
+        ca_cert = Path(ssl_config.certificates.stagingDir, CA_CERT_FILE_NAME)
+        write_data_to_file(
+            ca_cert, ssl_config.certificates.caCertValue.get_secret_value()
+        )
+    if ssl_config.certificates.clientCertValue:
+        client_cert = Path(ssl_config.certificates.stagingDir, CLIENT_CERT_FILE_NAME)
+        write_data_to_file(
+            client_cert,
+            ssl_config.certificates.clientCertValue.get_secret_value(),
+        )
+    if ssl_config.certificates.privateKeyValue:
+        private_key = Path(ssl_config.certificates.stagingDir, KEY_CERT_FILE_NAME)
+        write_data_to_file(
+            private_key,
+            ssl_config.certificates.privateKeyValue.get_secret_value(),
+        )
+    return ca_cert, client_cert, private_key
+
+
+def _handle_ssl_context_by_path(ssl_config: SslConfig):
+    ca_cert = False
+    if ssl_config.certificates.caCertPath:
+        ca_cert = ssl_config.certificates.caCertPath
+    client_cert = ssl_config.certificates.clientCertPath
+    private_key = ssl_config.certificates.privateKeyPath
+    return ca_cert, client_cert, private_key
+
+
+def get_connection(connection: OpenSearchConnection) -> OpenSearch:
+    """
+    Create OpenSearch connection supporting Basic and AWS IAM authentication.
+    """
+    basic_auth = None
+    aws_auth = None
+    verify_ssl = False
+    ssl_show_warn = False
+    ca_cert = False
+    client_cert = None
+    private_key = None
+
+    if connection.verifySSL == VerifySSL.validate:
+        verify_ssl = True
+    elif connection.verifySSL == VerifySSL.ignore:
+        ssl_show_warn = True
+
+    if connection.sslConfig and connection.sslConfig.certificates:
+        if isinstance(connection.sslConfig.certificates, SslCertificatesByValues):
+            ca_cert, client_cert, private_key = _handle_ssl_context_by_value(
+                ssl_config=connection.sslConfig
+            )
+        elif isinstance(connection.sslConfig.certificates, SslCertificatesByPath):
+            ca_cert, client_cert, private_key = _handle_ssl_context_by_path(
+                ssl_config=connection.sslConfig
+            )
+
+    # Check for Basic Authentication
+    if (
+        isinstance(connection.authType, BasicAuthentication)
+        and connection.authType.username
+    ):
+        basic_auth = (
+            connection.authType.username,
+            connection.authType.password.get_secret_value()
+            if connection.authType.password
+            else None,
+        )
+
+    # Check for AWS IAM Authentication
+    if isinstance(connection.authType, AWSCredentials):
+        aws_access_key = (
+            connection.authType.awsAccessKeyId.get_secret_value()
+            if connection.authType.awsAccessKeyId
+            else None
+        )
+        aws_secret_key = (
+            connection.authType.awsSecretAccessKey.get_secret_value()
+            if connection.authType.awsSecretAccessKey
+            else None
+        )
+        aws_region = connection.authType.awsRegion  # Region as a plain string
+        aws_session_token = (
+            connection.authType.awsSessionToken.get_secret_value()
+            if hasattr(connection.authType, "awsSessionToken")
+            and connection.authType.awsSessionToken
+            else None
+        )
+        aws_auth = AWS4Auth(
+            aws_access_key,
+            aws_secret_key,
+            aws_region,
+            "es",
+            session_token=aws_session_token,
+        )
+
+    if not connection.connectionArguments:
+        connection.connectionArguments = init_empty_connection_arguments()
+
+    # Determine the http_auth based on the available authentication method.
+    # AWS IAM takes precedence, followed by Basic Authentication.
+    http_auth = aws_auth if aws_auth else basic_auth
+
+    return OpenSearch(
+        clean_uri(str(connection.hostPort)),
+        http_auth=http_auth,
+        verify_certs=verify_ssl,
+        ssl_show_warn=ssl_show_warn,
+        ca_certs=ca_cert,
+        client_cert=client_cert,
+        client_key=private_key,
+        connection_class=RequestsHttpConnection,  # Use RequestsHttpConnection for AWS auth support.
+        **connection.connectionArguments.root,
+    )
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: OpenSearch,
+    service_connection: OpenSearchConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
+    """
+    Test connection for OpenSearch. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow.
+    """
+
+    def test_get_search_indexes():
+        client.indices.get_alias(expand_wildcards="open")
+
+    test_fn = {
+        "CheckAccess": client.info,
+        "GetSearchIndexes": test_get_search_indexes,
+    }
+
+    return test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
+    )

metadata/ingestion/source/search/opensearch/metadata.py

@@ -0,0 +1,270 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+"""
+OpenSearch source to extract metadata
+"""
+import shutil
+import traceback
+from pathlib import Path
+from typing import Any, Iterable, Optional
+
+from opensearchpy import OpenSearch
+
+from metadata.generated.schema.api.data.createSearchIndex import (
+    CreateSearchIndexRequest,
+)
+from metadata.generated.schema.entity.data.searchIndex import (
+    IndexType,
+    SearchIndex,
+    SearchIndexSampleData,
+)
+from metadata.generated.schema.entity.services.connections.search.openSearchConnection import (
+    OpenSearchConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.steps import InvalidSourceException, Source
+from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.search.opensearch.parser import parse_os_index_mapping
+from metadata.ingestion.source.search.search_service import SearchServiceSource
+from metadata.utils import fqn
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+WILDCARD_SEARCH = "*"
+
+
+class OpensearchSource(SearchServiceSource):
+    """
+    Implements the necessary methods to extract
+    Search Index metadata from OpenSearch.
+    Excludes system indexes (indexes whose names start with a dot).
+    """
+
+    def __init__(self, config: Source, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.client: OpenSearch = self.connection
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """
+        Create an instance of OpensearchSource.
+
+        Args:
+            config_dict: The configuration dictionary.
+            metadata: An instance of OpenMetadata.
+            pipeline_name: Optional pipeline name.
+
+        Returns:
+            An instance of OpensearchSource.
+        """
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: OpenSearchConnection = config.serviceConnection.root.config
+        if not isinstance(connection, OpenSearchConnection):
+            raise InvalidSourceException(
+                f"Expected OpenSearchConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def _is_system_index(self, index_name: str) -> bool:
+        """
+        Determine if the given index is a system index.
+        By default, system indexes start with a dot ('.').
+
+        Args:
+            index_name: The name of the index.
+
+        Returns:
+            True if it is a system index, False otherwise.
+        """
+        return index_name.startswith(".")
+
+    def get_search_index_list(self) -> Iterable[dict]:
+        """
+        Get a list of all search indices from OpenSearch, excluding system indexes.
+
+        Returns:
+            Iterable of dictionaries containing index details.
+        """
+        index_list = self.client.indices.get_alias(expand_wildcards="open") or {}
+        for index in index_list.keys():
+            if self._is_system_index(index):
+                logger.debug("Skipping system index: %s", index)
+                continue
+            yield self.client.indices.get(index=str(index))
+
+    def get_search_index_name(self, search_index_details: dict) -> Optional[str]:
+        """
+        Get the search index name.
+
+        Args:
+            search_index_details: Dictionary containing index details.
+
+        Returns:
+            The index name if available, else None.
+        """
+        if search_index_details and len(search_index_details) == 1:
+            return list(search_index_details.keys())[0]
+        return None
+
+    def yield_search_index(
+        self, search_index_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """
+        Yield Search Index entities.
+
+        Args:
+            search_index_details: The index details containing mappings and settings.
+
+        Yields:
+            Either wrapped CreateSearchIndexRequest objects.
+        """
+        index_name = self.get_search_index_name(search_index_details)
+        if index_name:
+            search_index_request = CreateSearchIndexRequest(
+                name=EntityName(index_name),
+                displayName=index_name,
+                searchIndexSettings=search_index_details.get(index_name, {}).get(
+                    "settings", {}
+                ),
+                service=FullyQualifiedEntityName(self.context.get().search_service),
+                fields=parse_os_index_mapping(
+                    search_index_details.get(index_name, {}).get("mappings")
+                ),
+                indexType=IndexType.Index,
+            )
+            yield Either(right=search_index_request)
+            self.register_record(search_index_request=search_index_request)
+
+    def yield_search_index_sample_data(
+        self, search_index_details: Any
+    ) -> Iterable[Either[OMetaIndexSampleData]]:
+        """
+        Yield sample data for the search index entity.
+
+        Args:
+            search_index_details: The index details used to query sample data.
+
+        Yields:
+            Either wrapped OMetaIndexSampleData objects.
+        """
+        if self.source_config.includeSampleData and self.context.get().search_index:
+            sample_data = self.client.search(
+                index=self.context.get().search_index,
+                q=WILDCARD_SEARCH,
+                size=self.source_config.sampleSize,
+                request_timeout=self.service_connection.connectionTimeoutSecs,
+            )
+
+            search_index_fqn = fqn.build(
+                metadata=self.metadata,
+                entity_type=SearchIndex,
+                service_name=self.context.get().search_service,
+                search_index_name=self.context.get().search_index,
+            )
+            search_index_entity = self.metadata.get_by_name(
+                entity=SearchIndex, fqn=search_index_fqn
+            )
+
+            yield Either(
+                right=OMetaIndexSampleData(
+                    entity=search_index_entity,
+                    data=SearchIndexSampleData(
+                        messages=[
+                            str(message)
+                            for message in sample_data.get("hits", {}).get("hits", [])
+                        ]
+                    ),
+                )
+            )
+
+    def get_search_index_template_list(self) -> Iterable[dict]:
+        """
+        Get a list of all search index templates from OpenSearch.
+
+        Returns:
+            Iterable of dictionaries containing template details.
+        """
+        yield from self.client.indices.get_index_template().get("index_templates", [])
+
+    def get_search_index_template_name(
+        self, search_index_template_details: dict
+    ) -> Optional[str]:
+        """
+        Get the search index template name.
+
+        Args:
+            search_index_template_details: Dictionary containing template details.
+
+        Returns:
+            The template name if available, else None.
+        """
+        return search_index_template_details and search_index_template_details["name"]
+
+    def yield_search_index_template(
+        self, search_index_template_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """
+        Yield Search Index Template entities.
+
+        Args:
+            search_index_template_details: The template details including mappings and settings.
+
+        Yields:
+            Either wrapped CreateSearchIndexRequest objects.
+        """
+        try:
+            if self.source_config.includeIndexTemplate:
+                index_name = self.get_search_index_template_name(
+                    search_index_template_details
+                )
+                index_template = search_index_template_details["index_template"]
+                if index_name:
+                    search_index_template_request = CreateSearchIndexRequest(
+                        name=EntityName(index_name),
+                        displayName=index_name,
+                        searchIndexSettings=index_template.get("template", {}).get(
+                            "settings", {}
+                        ),
+                        service=FullyQualifiedEntityName(
+                            self.context.get().search_service
+                        ),
+                        fields=parse_os_index_mapping(
+                            index_template.get("template", {}).get("mappings")
+                        ),
+                        indexType=IndexType.IndexTemplate,
+                        description=index_template.get("_meta", {}).get("description"),
+                    )
+                    yield Either(right=search_index_template_request)
+                    self.register_record(
+                        search_index_request=search_index_template_request
+                    )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Could not include index templates due to {exc}")
+
+    def close(self):
+        """
+        Clean up any temporary files and close the connection.
+        """
+        try:
+            if Path(self.service_connection.sslConfig.certificates.stagingDir).exists():
+                shutil.rmtree(self.service_connection.sslConfig.certificates.stagingDir)
+        except AttributeError:
+            pass
+        return super().close()