openmetadata-ingestion 1.4.7.2__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/ingestion/ometa/ometa_api.py

@@ -30,7 +30,6 @@ from metadata.generated.schema.type import basic
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.generated.schema.type.entityHistory import EntityVersionHistory
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.models.encoders import show_secrets_encoder
 from metadata.ingestion.ometa.auth_provider import OpenMetadataAuthenticationProvider
 from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
 from metadata.ingestion.ometa.mixins.custom_property_mixin import (
@@ -224,6 +223,8 @@ class OpenMetadata(
             class_name.lower()
             .replace("glossaryterm", "glossaryTerm")
             .replace("dashboarddatamodel", "dashboardDataModel")
+            .replace("apiendpoint", "apiEndpoint")
+            .replace("apicollection", "apiCollection")
             .replace("testsuite", "testSuite")
             .replace("testdefinition", "testDefinition")
             .replace("testcase", "testCase")
@@ -264,10 +265,10 @@ class OpenMetadata(
             )
 
         fn = getattr(self.client, method)
-        resp = fn(self.get_suffix(entity), data=data.json(encoder=show_secrets_encoder))
+        resp = fn(self.get_suffix(entity), data=data.model_dump_json())
         if not resp:
             raise EmptyPayloadException(
-                f"Got an empty response when trying to PUT to {self.get_suffix(entity)}, {data.json()}"
+                f"Got an empty response when trying to PUT to {self.get_suffix(entity)}, {data.model_dump_json()}"
             )
         return entity_class(**resp)
 
@@ -445,8 +446,7 @@ class OpenMetadata(
             params=params,
             skip_on_failure=skip_on_failure,
         )
-        for elem in entity_list.entities:
-            yield elem
+        yield from entity_list.entities
 
         after = entity_list.after
         while after:
@@ -458,8 +458,7 @@ class OpenMetadata(
                 after=after,
                 skip_on_failure=skip_on_failure,
             )
-            for elem in entity_list.entities:
-                yield elem
+            yield from entity_list.entities
             after = entity_list.after
 
     def list_versions(

metadata/ingestion/ometa/routes.py

@@ -21,6 +21,12 @@ from metadata.generated.schema.api.classification.createClassification import (
     CreateClassificationRequest,
 )
 from metadata.generated.schema.api.classification.createTag import CreateTagRequest
+from metadata.generated.schema.api.data.createAPICollection import (
+    CreateAPICollectionRequest,
+)
+from metadata.generated.schema.api.data.createAPIEndpoint import (
+    CreateAPIEndpointRequest,
+)
 from metadata.generated.schema.api.data.createChart import CreateChartRequest
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.api.data.createDashboard import CreateDashboardRequest
@@ -53,6 +59,9 @@ from metadata.generated.schema.api.domains.createDomain import CreateDomainReque
 from metadata.generated.schema.api.feed.createSuggestion import CreateSuggestionRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.api.policies.createPolicy import CreatePolicyRequest
+from metadata.generated.schema.api.services.createApiService import (
+    CreateApiServiceRequest,
+)
 from metadata.generated.schema.api.services.createDashboardService import (
     CreateDashboardServiceRequest,
 )
@@ -106,6 +115,8 @@ from metadata.generated.schema.entity.classification.classification import (
     Classification,
 )
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.data.apiCollection import APICollection
+from metadata.generated.schema.entity.data.apiEndpoint import APIEndpoint
 from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
@@ -127,6 +138,7 @@ from metadata.generated.schema.entity.domains.dataProduct import DataProduct
 from metadata.generated.schema.entity.domains.domain import Domain
 from metadata.generated.schema.entity.feed.suggestion import Suggestion
 from metadata.generated.schema.entity.policies.policy import Policy
+from metadata.generated.schema.entity.services.apiService import ApiService
 from metadata.generated.schema.entity.services.connections.testConnectionDefinition import (
     TestConnectionDefinition,
 )
@@ -179,6 +191,10 @@ ROUTES = {
     CreateSearchIndexRequest.__name__: "/searchIndexes",
     StoredProcedure.__name__: "/storedProcedures",
     CreateStoredProcedureRequest.__name__: "/storedProcedures",
+    APIEndpoint.__name__: "/apiEndpoints",
+    CreateAPIEndpointRequest.__name__: "/apiEndpoints",
+    APICollection.__name__: "/apiCollections",
+    CreateAPICollectionRequest.__name__: "/apiCollections",
     # Classifications
     Tag.__name__: "/tags",
     CreateTagRequest.__name__: "/tags",
@@ -205,6 +221,8 @@ ROUTES = {
     Workflow.__name__: "/automations/workflows",
     CreateWorkflowRequest.__name__: "/automations/workflows",
     # Services
+    ApiService.__name__: "/services/apiServices",
+    CreateApiServiceRequest.__name__: "/services/apiServices",
     DatabaseService.__name__: "/services/databaseServices",
     CreateDatabaseServiceRequest.__name__: "/services/databaseServices",
     DashboardService.__name__: "/services/dashboardServices",

metadata/ingestion/ometa/ttl_cache.py

@@ -0,0 +1,47 @@
+#  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.
+
+"""
+Simple dictionary implementation for keys with TTL
+"""
+from datetime import datetime
+from typing import Dict
+
+
+class TTLCache:
+    """
+    class to handle ttl cache
+    """
+
+    def __init__(self, ttl: int):
+        self._ttl = ttl
+        # The key will be the object, and the value the created time to check the TTL
+        self._cache: Dict[str, int] = {}
+
+    @staticmethod
+    def _now() -> int:
+        return int(datetime.now().timestamp())
+
+    def __contains__(self, item) -> bool:
+        if item in self._cache:
+            created_at = self._cache[item]
+            if self._now() - created_at > self._ttl:
+                self.delete(item)
+                return False
+            return True
+        return False
+
+    def add(self, value: str):
+        if value not in self._cache:
+            self._cache[value] = self._now()
+
+    def delete(self, key):
+        self._cache.pop(key, None)

metadata/ingestion/ometa/utils.py

@@ -68,9 +68,9 @@ def model_str(arg: Any) -> str:
     Default model stringifying method.
 
     Some elements such as FQN, EntityName, UUID
-    have the actual value under the pydantic base __root__
+    have the actual value under the pydantic base root
     """
-    if hasattr(arg, "__root__"):
-        return str(arg.__root__)
+    if hasattr(arg, "root"):
+        return str(arg.root)
 
     return str(arg)

metadata/ingestion/processor/query_parser.py

@@ -43,7 +43,7 @@ def parse_sql_statement(record: TableQuery, dialect: Dialect) -> Optional[Parsed
 
     start_time = record.analysisDate
     if isinstance(start_time, DateTime):
-        start_date = start_time.__root__.date()
+        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()))
@@ -62,7 +62,7 @@ def parse_sql_statement(record: TableQuery, dialect: Dialect) -> Optional[Parsed
         query_type=record.query_type,
         exclude_usage=record.exclude_usage,
         userName=record.userName,
-        date=start_time,
+        date=str(start_time),
         serviceName=record.serviceName,
         duration=record.duration,
     )
@@ -96,7 +96,7 @@ class QueryParserProcessor(Processor):
         pipeline_name: Optional[str] = None,
         **kwargs,
     ):
-        config = ConfigModel.parse_obj(config_dict)
+        config = ConfigModel.model_validate(config_dict)
         connection_type = kwargs.pop("connection_type", "")
         return cls(config, metadata, connection_type)
 

metadata/ingestion/sink/file.py

@@ -53,14 +53,14 @@ class FileSink(Sink):
     def create(
         cls, config_dict: dict, _: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config = FileSinkConfig.parse_obj(config_dict)
+        config = FileSinkConfig.model_validate(config_dict)
         return cls(config)
 
     def _run(self, record: Entity, *_, **__) -> Either[str]:
         if self.wrote_something:
             self.file.write(",\n")
 
-        self.file.write(record.json())
+        self.file.write(record.model_dump_json())
         self.wrote_something = True
         return Either(right=get_log_name(record))
 

metadata/ingestion/sink/metadata_rest.py

@@ -21,7 +21,6 @@ from pydantic import BaseModel
 from requests.exceptions import HTTPError
 
 from metadata.config.common import ConfigModel
-from metadata.data_insight.source.metadata import DataInsightRecord
 from metadata.data_quality.api.models import TestCaseResultResponse, TestCaseResults
 from metadata.generated.schema.analytics.reportData import ReportData
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
@@ -51,6 +50,7 @@ from metadata.generated.schema.tests.testCaseResolutionStatus import (
     TestCaseResolutionStatus,
 )
 from metadata.generated.schema.tests.testSuite import TestSuite
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.schema import Topic
 from metadata.ingestion.api.models import Either, Entity, StackTraceError
 from metadata.ingestion.api.steps import Sink
@@ -59,6 +59,7 @@ from metadata.ingestion.models.data_insight import OMetaDataInsightSample
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.life_cycle import OMetaLifeCycleData
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.ometa_topic_data import OMetaTopicSampleData
 from metadata.ingestion.models.patch_request import (
     ALLOWED_COMMON_PATCH_FIELDS,
@@ -123,7 +124,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         metadata: OpenMetadata,
         pipeline_name: Optional[str] = None,
     ):
-        config = MetadataRestSinkConfig.parse_obj(config_dict)
+        config = MetadataRestSinkConfig.model_validate(config_dict)
         return cls(config, metadata)
 
     @property
@@ -186,6 +187,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             allowed_fields=ALLOWED_COMMON_PATCH_FIELDS,
             restrict_update_fields=RESTRICT_UPDATE_LIST,
             array_entity_fields=ARRAY_ENTITY_FIELDS,
+            override_metadata=record.override_metadata,
         )
         patched_entity = PatchedEntity(new_entity=entity) if entity else None
         return Either(right=patched_entity)
@@ -249,13 +251,51 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         created_lineage = self.metadata.add_lineage(add_lineage, check_patch=True)
         return Either(right=created_lineage["entity"]["fullyQualifiedName"])
 
+    @_run_dispatch.register
+    def write_override_lineage(
+        self, add_lineage: OMetaLineageRequest
+    ) -> Either[Dict[str, Any]]:
+        """
+        Writes the override lineage for the given lineage request.
+
+        Args:
+            add_lineage (OMetaLineageRequest): The lineage request containing the override lineage information.
+
+        Returns:
+            Either[Dict[str, Any]]: The result of the dispatch operation.
+        """
+        if (
+            add_lineage.override_lineage is True
+            and add_lineage.lineage_request.edge.lineageDetails
+            and add_lineage.lineage_request.edge.lineageDetails.source
+        ):
+            if (
+                add_lineage.lineage_request.edge.lineageDetails.pipeline
+                and add_lineage.lineage_request.edge.lineageDetails.source
+                == LineageSource.PipelineLineage
+            ):
+                self.metadata.delete_lineage_by_source(
+                    entity_type="pipeline",
+                    entity_id=str(
+                        add_lineage.lineage_request.edge.lineageDetails.pipeline.id.root
+                    ),
+                    source=add_lineage.lineage_request.edge.lineageDetails.source.value,
+                )
+            else:
+                self.metadata.delete_lineage_by_source(
+                    entity_type=add_lineage.lineage_request.edge.toEntity.type,
+                    entity_id=str(add_lineage.lineage_request.edge.toEntity.id.root),
+                    source=add_lineage.lineage_request.edge.lineageDetails.source.value,
+                )
+        return self._run_dispatch(add_lineage.lineage_request)
+
     def _create_role(self, create_role: CreateRoleRequest) -> Optional[Role]:
         """
         Internal helper method for write_user
         """
         try:
             role = self.metadata.create_or_update(create_role)
-            self.role_entities[role.name] = str(role.id.__root__)
+            self.role_entities[role.name] = str(role.id.root)
             return role
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -269,7 +309,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         """
         try:
             team = self.metadata.create_or_update(create_team)
-            self.team_entities[team.name.__root__] = str(team.id.__root__)
+            self.team_entities[team.name.root] = str(team.id.root)
             return team
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -292,7 +332,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             for role in record.roles:
                 try:
                     role_entity = self.metadata.get_by_name(
-                        entity=Role, fqn=str(role.name.__root__)
+                        entity=Role, fqn=str(role.name.root)
                     )
                 except APIError:
                     role_entity = self._create_role(role)
@@ -307,18 +347,16 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             for team in record.teams:
                 try:
                     team_entity = self.metadata.get_by_name(
-                        entity=Team, fqn=str(team.name.__root__)
+                        entity=Team, fqn=str(team.name.root)
                     )
                     if not team_entity:
                         raise APIError(
-                            error={
-                                "message": f"Creating a new team {team.name.__root__}"
-                            }
+                            error={"message": f"Creating a new team {team.name.root}"}
                         )
-                    team_ids.append(team_entity.id.__root__)
+                    team_ids.append(team_entity.id.root)
                 except APIError:
                     team_entity = self._create_team(team)
-                    team_ids.append(team_entity.id.__root__)
+                    team_ids.append(team_entity.id.root)
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
                     logger.warning(f"Unexpected error writing team [{team}]: {exc}")
@@ -326,7 +364,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             team_ids = None
 
         # Update user data with the new Role and Team IDs
-        user_profile = record.user.dict(exclude_unset=True)
+        user_profile = record.user.model_dump(exclude_unset=True)
         user_profile["roles"] = role_ids
         user_profile["teams"] = team_ids
         metadata_user = CreateUserRequest(**user_profile)
@@ -421,10 +459,10 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         """Write the test case result"""
         res = self.metadata.add_test_case_results(
             test_results=record.testCaseResult,
-            test_case_fqn=record.testCase.fullyQualifiedName.__root__,
+            test_case_fqn=record.testCase.fullyQualifiedName.root,
         )
         logger.debug(
-            f"Successfully ingested test case results for test case {record.testCase.name.__root__}"
+            f"Successfully ingested test case results for test case {record.testCase.name.root}"
         )
         return Either(right=res)
 
@@ -449,20 +487,12 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         )
         return Either(right=record.record)
 
-    @_run_dispatch.register
-    def write_data_insight(self, record: DataInsightRecord) -> Either[ReportData]:
-        """
-        Use the /dataQuality/testCases endpoint to ingest sample test suite
-        """
-        self.metadata.add_data_insight_report_data(record.data)
-        return Either(left=None, right=record.data)
-
     @_run_dispatch.register
     def write_data_insight_kpi(self, record: KpiResult) -> Either[KpiResult]:
         """
         Use the /dataQuality/testCases endpoint to ingest sample test suite
         """
-        self.metadata.add_kpi_result(fqn=record.kpiFqn.__root__, record=record)
+        self.metadata.add_kpi_result(fqn=record.kpiFqn.root, record=record)
         return Either(left=None, right=record)
 
     @_run_dispatch.register
@@ -532,7 +562,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             profile_request=record.profile,
         )
         logger.debug(
-            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.__root__}"
+            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
         )
 
         if record.sample_data:
@@ -542,13 +572,13 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             if not table_data:
                 self.status.failed(
                     StackTraceError(
-                        name=table.fullyQualifiedName.__root__,
+                        name=table.fullyQualifiedName.root,
                         error="Error trying to ingest sample data for table",
                     )
                 )
             else:
                 logger.debug(
-                    f"Successfully ingested sample data for {record.table.fullyQualifiedName.__root__}"
+                    f"Successfully ingested sample data for {record.table.fullyQualifiedName.root}"
                 )
 
         if record.column_tags:
@@ -558,13 +588,13 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             if not patched:
                 self.status.failed(
                     StackTraceError(
-                        name=table.fullyQualifiedName.__root__,
+                        name=table.fullyQualifiedName.root,
                         error="Error patching tags for table",
                     )
                 )
             else:
                 logger.debug(
-                    f"Successfully patched tag {record.column_tags} for {record.table.fullyQualifiedName.__root__}"
+                    f"Successfully patched tag {record.column_tags} for {record.table.fullyQualifiedName.root}"
                 )
 
         return Either(right=table)
@@ -586,7 +616,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         for result in record.test_results or []:
             self.metadata.add_test_case_results(
                 test_results=result.testCaseResult,
-                test_case_fqn=result.testCase.fullyQualifiedName.__root__,
+                test_case_fqn=result.testCase.fullyQualifiedName.root,
             )
             self.status.scanned(result)
 

metadata/ingestion/source/connections.py

@@ -13,11 +13,16 @@
 Main entrypoints to create and test connections
 for any source.
 """
+import traceback
 from typing import Any, Callable
 
 from pydantic import BaseModel
+from sqlalchemy.engine import Engine
 
 from metadata.utils.importer import import_connection_fn
+from metadata.utils.logger import cli_logger
+
+logger = cli_logger()
 
 GET_CONNECTION_FN_NAME = "get_connection"
 TEST_CONNECTION_FN_NAME = "test_connection"
@@ -47,3 +52,17 @@ def get_connection(connection: BaseModel) -> Any:
     a service connection pydantic model
     """
     return get_connection_fn(connection)(connection)
+
+
+def kill_active_connections(engine: Engine):
+    """
+    Method to kill the active connections
+    as well as idle connections in the engine
+    """
+    try:
+        active_conn = engine.pool.checkedout() + engine.pool.checkedin()
+        if active_conn:
+            engine.dispose()
+    except Exception as exc:
+        logger.warning(f"Error Killing the active connections {exc}")
+        logger.debug(traceback.format_exc())