openmetadata-ingestion 1.5.13.1__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/ingestion/ometa/mixins/data_insight_mixin.py

@@ -31,6 +31,7 @@ from metadata.generated.schema.dataInsight.dataInsightChartResult import (
 )
 from metadata.generated.schema.dataInsight.kpi.basic import KpiResult
 from metadata.generated.schema.dataInsight.kpi.kpi import Kpi
+from metadata.ingestion.ometa.utils import quote
 
 
 class DataInsightMixin:
@@ -58,7 +59,7 @@ class DataInsightMixin:
             record (ReportData): report data
         """
 
-        resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.model_dump_json())
+        resp = self.client.put(f"/kpi/{quote(fqn)}/kpiResult", record.model_dump_json())
 
         return resp
 
@@ -143,7 +144,7 @@ class DataInsightMixin:
         params = {"startTs": start_ts, "endTs": end_ts}
 
         resp = self.client.get(
-            f"/kpi/{fqn}/kpiResult",
+            f"/kpi/{quote(fqn)}/kpiResult",
             params,
         )
 

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/mixins/ingestion_pipeline_mixin.py

@@ -22,6 +22,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
 )
 from metadata.ingestion.api.parser import parse_ingestion_pipeline_config_gracefully
 from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.utils import quote
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -46,7 +47,7 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status: Pipeline Status data to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=pipeline_status.model_dump_json(),
         )
         logger.debug(
@@ -64,7 +65,8 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status_run_id: Pipeline Status run id
         """
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus/{pipeline_status_run_id}"
+            f"{self.get_suffix(IngestionPipeline)}/"
+            f"{quote(ingestion_pipeline_fqn)}/pipelineStatus/{pipeline_status_run_id}"
         )
         if resp:
             return PipelineStatus(**resp)
@@ -99,7 +101,7 @@ class OMetaIngestionPipelineMixin:
         params = {"startTs": start_ts, "endTs": end_ts}
 
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=params,
         )
 

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -19,7 +19,7 @@ from metadata.generated.schema.entity.feed.suggestion import Suggestion, Suggest
 from metadata.generated.schema.type import basic
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.ingestion.ometa.client import REST
-from metadata.ingestion.ometa.utils import model_str
+from metadata.ingestion.ometa.utils import model_str, quote
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -65,7 +65,7 @@ class OMetaSuggestionsMixin:
         self.client.put(
             f"{self.get_suffix(Suggestion)}/accept-all?"
             f"userId={model_str(user_id)}&"
-            f"entityFQN={model_str(fqn)}&"
+            f"entityFQN={quote(fqn)}&"
             f"suggestionType={suggestion_type.value}",
         )
 
@@ -79,6 +79,6 @@ class OMetaSuggestionsMixin:
         self.client.put(
             f"{self.get_suffix(Suggestion)}/reject-all?"
             f"userId={model_str(user_id)}&"
-            f"entityFQN={model_str(fqn)}&"
+            f"entityFQN={quote(fqn)}&"
             f"suggestionType={suggestion_type.value}",
         )

metadata/ingestion/ometa/mixins/table_mixin.py

@@ -257,7 +257,7 @@ class OMetaTableMixin:
         profile_type_url = profile_type.__name__[0].lower() + profile_type.__name__[1:]
 
         resp = self.client.get(
-            f"{self.get_suffix(Table)}/{fqn}/{profile_type_url}?limit={limit}{url_after}",
+            f"{self.get_suffix(Table)}/{quote(fqn)}/{profile_type_url}?limit={limit}{url_after}",
             data={"startTs": start_ts, "endTs": end_ts},
         )
 

metadata/ingestion/ometa/models.py

@@ -30,3 +30,4 @@ class EntityList(BaseModel, Generic[T]):
     entities: List[T]
     total: int
     after: Optional[str] = None
+    before: Optional[str] = None

metadata/ingestion/ometa/ometa_api.py

@@ -378,12 +378,13 @@ class OpenMetadata(
         logger.debug("Cannot find the Entity %s", fqn)
         return None
 
-    # pylint: disable=too-many-locals
+    # pylint: disable=too-many-locals, too-many-arguments
     def list_entities(
         self,
         entity: Type[T],
         fields: Optional[List[str]] = None,
         after: Optional[str] = None,
+        before: Optional[str] = None,
         limit: int = 100,
         params: Optional[Dict[str, str]] = None,
         skip_on_failure: bool = False,
@@ -395,9 +396,10 @@ class OpenMetadata(
         suffix = self.get_suffix(entity)
         url_limit = f"?limit={limit}"
         url_after = f"&after={after}" if after else ""
+        url_before = f"&before={before}" if before else ""
         url_fields = f"&fields={','.join(fields)}" if fields else ""
         resp = self.client.get(
-            path=f"{suffix}{url_limit}{url_after}{url_fields}", data=params
+            path=f"{suffix}{url_limit}{url_after}{url_before}{url_fields}", data=params
         )
 
         if self._use_raw_data:
@@ -421,7 +423,8 @@ class OpenMetadata(
 
         total = resp["paging"]["total"]
         after = resp["paging"]["after"] if "after" in resp["paging"] else None
-        return EntityList(entities=entities, total=total, after=after)
+        before = resp["paging"]["before"] if "before" in resp["paging"] else None
+        return EntityList(entities=entities, total=total, after=after, before=before)
 
     def list_all_entities(
         self,

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"""