openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/ometa/routes.py

@@ -42,6 +42,7 @@ from metadata.generated.schema.api.data.createGlossary import CreateGlossaryRequ
 from metadata.generated.schema.api.data.createGlossaryTerm import (
     CreateGlossaryTermRequest,
 )
+from metadata.generated.schema.api.data.createMetric import CreateMetricRequest
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -126,7 +127,7 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.glossary import Glossary
 from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
-from metadata.generated.schema.entity.data.metrics import Metrics
+from metadata.generated.schema.entity.data.metric import Metric
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
@@ -181,7 +182,8 @@ ROUTES = {
     CreateTableRequest.__name__: "/tables",
     Topic.__name__: "/topics",
     CreateTopicRequest.__name__: "/topics",
-    Metrics.__name__: "/metrics",
+    Metric.__name__: "/metrics",
+    CreateMetricRequest.__name__: "/metrics",
     AddLineageRequest.__name__: "/lineage",
     Report.__name__: "/reports",
     Query.__name__: "/queries",

metadata/ingestion/processor/query_parser.py

@@ -26,7 +26,7 @@ from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Diale
 from metadata.ingestion.lineage.parser import LineageParser
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
+from metadata.utils.time_utils import datetime_to_timestamp
 
 logger = ingestion_logger()
 
@@ -46,7 +46,7 @@ def parse_sql_statement(record: TableQuery, dialect: Dialect) -> Optional[Parsed
         start_date = start_time.root.date()
         start_time = datetime.datetime.strptime(str(start_date.isoformat()), "%Y-%m-%d")
 
-    start_time = convert_timestamp_to_milliseconds(int(start_time.timestamp()))
+    start_time = datetime_to_timestamp(start_time, milliseconds=True)
 
     lineage_parser = LineageParser(record.query, dialect=dialect)
 

metadata/ingestion/sink/metadata_rest.py

@@ -84,6 +84,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardUsage
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.profiler.api.models import ProfilerResponse
+from metadata.sampler.models import SamplerResponse
 from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import get_log_name, ingestion_logger
 
@@ -549,22 +550,8 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         )
 
     @_run_dispatch.register
-    def write_profiler_response(self, record: ProfilerResponse) -> Either[Table]:
-        """Cleanup "`" character in columns and ingest"""
-        column_profile = record.profile.columnProfile
-        for column in column_profile:
-            column.name = column.name.replace("`", "")
-
-        record.profile.columnProfile = column_profile
-
-        table = self.metadata.ingest_profile_data(
-            table=record.table,
-            profile_request=record.profile,
-        )
-        logger.debug(
-            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
-        )
-
+    def write_sampler_response(self, record: SamplerResponse) -> Either[Table]:
+        """Ingest the sample data - if needed - and the PII tags"""
         if record.sample_data and record.sample_data.store:
             table_data = self.metadata.ingest_table_sample_data(
                 table=record.table, sample_data=record.sample_data.data
@@ -572,7 +559,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             if not table_data:
                 self.status.failed(
                     StackTraceError(
-                        name=table.fullyQualifiedName.root,
+                        name=record.table.fullyQualifiedName.root,
                         error="Error trying to ingest sample data for table",
                     )
                 )
@@ -587,7 +574,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             )
             if not patched:
                 self.status.warning(
-                    key=table.fullyQualifiedName.root,
+                    key=record.table.fullyQualifiedName.root,
                     reason="Error patching tags for table",
                 )
             else:
@@ -595,6 +582,24 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
                     f"Successfully patched tag {record.column_tags} for {record.table.fullyQualifiedName.root}"
                 )
 
+        return Either(right=record.table)
+
+    @_run_dispatch.register
+    def write_profiler_response(self, record: ProfilerResponse) -> Either[Table]:
+        """Cleanup "`" character in columns and ingest"""
+        column_profile = record.profile.columnProfile
+        for column in column_profile:
+            column.name = column.name.replace("`", "")
+
+        record.profile.columnProfile = column_profile
+
+        table = self.metadata.ingest_profile_data(
+            table=record.table,
+            profile_request=record.profile,
+        )
+        logger.debug(
+            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
+        )
         return Either(right=table)
 
     @_run_dispatch.register

metadata/ingestion/source/api/api_service.py

@@ -0,0 +1,212 @@
+#  Copyright 2024 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.
+"""
+Base class for ingesting api services
+"""
+from abc import ABC, abstractmethod
+from typing import Any, Iterable, Set
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from metadata.generated.schema.api.data.createAPICollection import (
+    CreateAPICollectionRequest,
+)
+from metadata.generated.schema.api.data.createAPIEndpoint import (
+    CreateAPIEndpointRequest,
+)
+from metadata.generated.schema.entity.data.apiCollection import APICollection
+from metadata.generated.schema.entity.data.apiEndpoint import APIEndpoint
+from metadata.generated.schema.entity.services.apiService import (
+    ApiConnection,
+    ApiService,
+)
+from metadata.generated.schema.metadataIngestion.apiServiceMetadataPipeline import (
+    ApiServiceMetadataPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.delete import delete_entity_from_source
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.steps import Source
+from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
+from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.topology import (
+    NodeStage,
+    ServiceTopology,
+    TopologyContextManager,
+    TopologyNode,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.utils import fqn
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class ApiServiceTopology(ServiceTopology):
+    """
+    Defines the hierarchy in API Services.
+    service -> ApiCollection -> ApiEndpoint
+
+    We could have a topology validator. We can only consume
+    data that has been produced by any parent node.
+    """
+
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
+        producer="get_services",
+        stages=[
+            NodeStage(
+                type_=ApiService,
+                context="api_service",
+                processor="yield_create_request_api_service",
+                overwrite=False,
+                must_return=True,
+                cache_entities=True,
+            ),
+        ],
+        children=["api_collection"],
+        post_process=["mark_api_collections_as_deleted"],
+    )
+    api_collection: Annotated[
+        TopologyNode, Field(description="API Collection Processing Node")
+    ] = TopologyNode(
+        producer="get_api_collections",
+        stages=[
+            NodeStage(
+                type_=APICollection,
+                context="api_collections",
+                processor="yield_api_collection",
+                consumer=["api_service"],
+                use_cache=True,
+            ),
+            NodeStage(
+                type_=APIEndpoint,
+                context="api_endpoints",
+                processor="yield_api_endpoint",
+                consumer=["api_service"],
+                use_cache=True,
+            ),
+        ],
+    )
+
+
+class ApiServiceSource(TopologyRunnerMixin, Source, ABC):
+    """
+    Base class for API services.
+    It implements the topology and context
+    """
+
+    source_config: ApiServiceMetadataPipeline
+    config: WorkflowSource
+    # Big union of types we want to fetch dynamically
+    service_connection: ApiConnection.model_fields["config"].annotation
+
+    topology = ApiServiceTopology()
+    context = TopologyContextManager(topology)
+    api_collection_source_state: Set = set()
+    api_endpoint_source_state: Set = set()
+
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
+        super().__init__()
+        self.config = config
+        self.metadata = metadata
+        self.service_connection = self.config.serviceConnection.root.config
+        self.source_config: ApiServiceMetadataPipeline = self.config.sourceConfig.config
+        self.connection = get_connection(self.service_connection)
+
+        # Flag the connection for the test connection
+        self.connection_obj = self.connection
+        self.test_connection()
+
+        self.client = self.connection
+
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
+    def get_services(self) -> Iterable[WorkflowSource]:
+        yield self.config
+
+    def yield_create_request_api_service(self, config: WorkflowSource):
+        yield Either(
+            right=self.metadata.get_create_service_from_source(
+                entity=ApiService, config=config
+            )
+        )
+
+    @abstractmethod
+    def get_api_collections(self, *args, **kwargs) -> Iterable[Any]:
+        """
+        Method to list all collections to process.
+        Here is where filtering happens
+        """
+
+    @abstractmethod
+    def yield_api_collection(
+        self, *args, **kwargs
+    ) -> Iterable[Either[CreateAPICollectionRequest]]:
+        """Method to return api collection Entities"""
+
+    @abstractmethod
+    def yield_api_endpoint(
+        self, *args, **kwargs
+    ) -> Iterable[Either[CreateAPIEndpointRequest]]:
+        """Method to return api endpoint Entities"""
+
+    def close(self):
+        """By default, nothing to close"""
+
+    def test_connection(self) -> None:
+        test_connection_fn = get_test_connection_fn(self.service_connection)
+        result = test_connection_fn(
+            self.metadata, self.connection_obj, self.service_connection
+        )
+        raise_test_connection_exception(result)
+
+    def mark_api_collections_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
+        """Method to mark the api collection as deleted"""
+        if self.source_config.markDeletedApiCollections:
+            yield from delete_entity_from_source(
+                metadata=self.metadata,
+                entity_type=APICollection,
+                entity_source_state=self.api_collection_source_state,
+                mark_deleted_entity=self.source_config.markDeletedApiCollections,
+                params={"service": self.context.get().api_service},
+            )
+
+    def register_record(self, collection_request: CreateAPICollectionRequest) -> None:
+        """
+        Mark the api collection record as scanned and update
+        the api_collection_source_state
+        """
+        api_collection_fqn = fqn.build(
+            self.metadata,
+            entity_type=APICollection,
+            service_name=collection_request.service.root,
+            api_collection_name=collection_request.name.root,
+        )
+
+        self.api_collection_source_state.add(api_collection_fqn)
+
+    def prepare(self):
+        """By default, nothing to prepare"""

metadata/ingestion/source/api/rest/connection.py

@@ -0,0 +1,93 @@
+#  Copyright 2024 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
+"""
+from typing import Optional
+
+import requests
+from requests.models import Response
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.api.restConnection import (
+    RestConnection,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+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
+
+
+class SchemaURLError(Exception):
+    """
+    Class to indicate schema url is invalid
+    """
+
+
+class InvalidOpenAPISchemaError(Exception):
+    """
+    Class to indicate openapi schema is invalid
+    """
+
+
+def get_connection(connection: RestConnection) -> Response:
+    """
+    Create connection
+    """
+    if connection.token:
+        headers = {"Authorization": f"Bearer {connection.token.get_secret_value()}"}
+        return requests.get(connection.openAPISchemaURL, headers=headers)
+    return requests.get(connection.openAPISchemaURL)
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: Response,
+    service_connection: RestConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+
+    def custom_url_exec():
+        if (
+            "application/json" in client.headers.get("content-type")
+            and client.status_code == 200
+        ):
+            return []
+        raise SchemaURLError(
+            "Failed to parse JSON schema url. Please check if provided url is valid JSON schema."
+        )
+
+    def custom_schema_exec():
+        if client.json().get("openapi") is not None:
+            return []
+        raise InvalidOpenAPISchemaError(
+            "Provided schema is not valid OpenAPI JSON schema"
+        )
+
+    test_fn = {"CheckURL": custom_url_exec, "CheckSchema": custom_schema_exec}
+
+    return test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
+    )