openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -24,9 +24,11 @@ from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.table import Table, TableType
 from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.ingestion.ometa.client import REST, APIError
 from metadata.ingestion.ometa.utils import quote
+from metadata.ingestion.source.models import TableView
 from metadata.utils.elasticsearch import ES_INDEX_MAP
 from metadata.utils.logger import ometa_logger
 
@@ -90,7 +92,7 @@ class ESMixin(Generic[T]):
 
     # sort_field needs to be unique for the pagination to work, so we can use the FQN
     paginate_query = (
-        "/search/query?q=&size={size}&deleted=false{filter}&index={index}"
+        "/search/query?q=&size={size}&deleted=false{filter}&index={index}{include_fields}"
         "&sort_field=fullyQualifiedName{after}"
     )
 
@@ -301,13 +303,19 @@ class ESMixin(Generic[T]):
             logger.warning(f"Unknown error extracting results from ES query [{err}]")
             return None
 
-    def paginate_es(
+    def _get_include_fields_query(self, fields: Optional[List[str]]) -> str:
+        """Get the include fields query"""
+        if fields:
+            return "&include_source_fields=" + "&include_source_fields=".join(fields)
+        return ""
+
+    def _paginate_es_internal(
         self,
         entity: Type[T],
         query_filter: Optional[str] = None,
         size: int = 100,
-        fields: Optional[List[str]] = None,
-    ) -> Iterator[T]:
+        include_fields: Optional[List[str]] = None,
+    ) -> Iterator[ESResponse]:
         """Paginate through the ES results, ignoring individual errors"""
         after: Optional[str] = None
         error_pages = 0
@@ -316,6 +324,7 @@ class ESMixin(Generic[T]):
             index=ES_INDEX_MAP[entity.__name__],
             filter="&query_filter=" + quote_plus(query_filter) if query_filter else "",
             size=size,
+            include_fields=self._get_include_fields_query(include_fields),
         )
         while True:
             query_string = query(
@@ -330,10 +339,7 @@ class ESMixin(Generic[T]):
                     continue
                 else:
                     break
-
-            yield from self._yield_hits_from_api(
-                response=response, entity=entity, fields=fields
-            )
+            yield response
 
             # Get next page
             last_hit = response.hits.hits[-1] if response.hits.hits else None
@@ -343,6 +349,18 @@ class ESMixin(Generic[T]):
 
             after = ",".join(last_hit.sort)
 
+    def paginate_es(
+        self,
+        entity: Type[T],
+        query_filter: Optional[str] = None,
+        size: int = 100,
+        fields: Optional[List[str]] = None,
+    ) -> Iterator[T]:
+        for response in self._paginate_es_internal(entity, query_filter, size):
+            yield from self._yield_hits_from_api(
+                response=response, entity=entity, fields=fields
+            )
+
     def _get_es_response(self, query_string: str) -> Optional[ESResponse]:
         """Get the Elasticsearch response"""
         try:
@@ -369,3 +387,43 @@ class ESMixin(Generic[T]):
                 logger.warning(
                     f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
                 )
+
+    def yield_es_view_def(
+        self,
+        service_name: str,
+    ) -> Iterable[TableView]:
+        """
+        Get the view definition from ES
+        """
+
+        from metadata.utils import fqn
+
+        query = {
+            "query": {
+                "bool": {
+                    "must": [
+                        {"term": {"service.name.keyword": service_name}},
+                        {"term": {"tableType": TableType.View.value}},
+                        {"term": {"deleted": False}},
+                        {"exists": {"field": "schemaDefinition"}},
+                    ]
+                }
+            }
+        }
+        query = json.dumps(query)
+        for response in self._paginate_es_internal(
+            entity=Table,
+            query_filter=query,
+            include_fields=["schemaDefinition", "fullyQualifiedName"],
+        ):
+            for hit in response.hits.hits:
+                _, database_name, schema_name, table_name = fqn.split(
+                    hit.source["fullyQualifiedName"]
+                )
+                yield TableView(
+                    view_definition=hit.source["schemaDefinition"],
+                    service_name=service_name,
+                    db_name=database_name,
+                    schema_name=schema_name,
+                    table_name=table_name,
+                )

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/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,
+    )