openmetadata-ingestion 1.5.7.1.dev0__py3-none-any.whl → 1.5.9.0__py3-none-any.whl

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -13,8 +13,13 @@ Mixin class containing Suggestions specific methods
 
 To be used by OpenMetadata class
 """
-from metadata.generated.schema.entity.feed.suggestion import Suggestion
+from typing import Union
+
+from metadata.generated.schema.entity.feed.suggestion import Suggestion, SuggestionType
+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.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -30,12 +35,50 @@ class OMetaSuggestionsMixin:
     client: REST
 
     def update_suggestion(self, suggestion: Suggestion) -> Suggestion:
-        """
-        Update an existing Suggestion with new fields
-        """
+        """Update an existing Suggestion with new fields"""
         resp = self.client.put(
             f"{self.get_suffix(Suggestion)}/{str(suggestion.root.id.root)}",
             data=suggestion.model_dump_json(),
         )
 
         return Suggestion(**resp)
+
+    def accept_suggestion(self, suggestion_id: Union[str, basic.Uuid]) -> None:
+        """Accept a given suggestion"""
+        self.client.put(
+            f"{self.get_suffix(Suggestion)}/{model_str(suggestion_id)}/accept",
+        )
+
+    def reject_suggestion(self, suggestion_id: Union[str, basic.Uuid]) -> None:
+        """Reject a given suggestion"""
+        self.client.put(
+            f"{self.get_suffix(Suggestion)}/{model_str(suggestion_id)}/reject",
+        )
+
+    def accept_all_suggestions(
+        self,
+        fqn: Union[str, FullyQualifiedEntityName],
+        user_id: Union[str, basic.Uuid],
+        suggestion_type: SuggestionType = SuggestionType.SuggestDescription,
+    ) -> None:
+        """Accept all suggestions"""
+        self.client.put(
+            f"{self.get_suffix(Suggestion)}/accept-all?"
+            f"userId={model_str(user_id)}&"
+            f"entityFQN={model_str(fqn)}&"
+            f"suggestionType={suggestion_type.value}",
+        )
+
+    def reject_all_suggestions(
+        self,
+        fqn: Union[str, FullyQualifiedEntityName],
+        user_id: Union[str, basic.Uuid],
+        suggestion_type: SuggestionType = SuggestionType.SuggestDescription,
+    ) -> None:
+        """Accept all suggestions"""
+        self.client.put(
+            f"{self.get_suffix(Suggestion)}/reject-all?"
+            f"userId={model_str(user_id)}&"
+            f"entityFQN={model_str(fqn)}&"
+            f"suggestionType={suggestion_type.value}",
+        )

metadata/ingestion/ometa/ometa_api.py

@@ -19,6 +19,7 @@ from typing import Dict, Generic, Iterable, List, Optional, Type, TypeVar, Union
 
 from pydantic import BaseModel
 
+from metadata.generated.schema.api.createBot import CreateBot
 from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
     CreateIngestionPipelineRequest,
 )
@@ -172,13 +173,16 @@ class OpenMetadata(
 
         return route
 
-    def get_module_path(self, entity: Type[T]) -> str:
+    def get_module_path(self, entity: Type[T]) -> Optional[str]:
         """
         Based on the entity, return the module path
         it is found inside generated
         """
         if issubclass(entity, CreateIngestionPipelineRequest):
             return "services.ingestionPipelines"
+        if issubclass(entity, CreateBot):
+            # Bots schemas don't live inside any subdirectory
+            return None
         return entity.__module__.split(".")[-2]
 
     def get_create_entity_type(self, entity: Type[T]) -> Type[C]:

metadata/ingestion/ometa/routes.py

@@ -21,6 +21,7 @@ from metadata.generated.schema.api.classification.createClassification import (
     CreateClassificationRequest,
 )
 from metadata.generated.schema.api.classification.createTag import CreateTagRequest
+from metadata.generated.schema.api.createBot import CreateBot
 from metadata.generated.schema.api.data.createAPICollection import (
     CreateAPICollectionRequest,
 )
@@ -211,7 +212,8 @@ ROUTES = {
     User.__name__: "/users",
     CreateUserRequest.__name__: "/users",
     AuthenticationMechanism.__name__: "/users/auth-mechanism",
-    Bot.__name__: "/bots",  # We won't allow bot creation from the client
+    Bot.__name__: "/bots",
+    CreateBot.__name__: "/bots",
     # Roles
     Role.__name__: "/roles",
     CreateRoleRequest.__name__: "/roles",

metadata/ingestion/source/dashboard/dashboard_service.py

@@ -55,6 +55,9 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either, Entity
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import C, TopologyRunnerMixin
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -554,7 +557,10 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
-        test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
+        result = test_connection_fn(
+            self.metadata, self.connection_obj, self.service_connection
+        )
+        raise_test_connection_exception(result)
 
     def prepare(self):
         """By default, nothing to prepare"""

metadata/ingestion/source/dashboard/domodashboard/connection.py

@@ -12,7 +12,6 @@
 """
 Source connection handler
 """
-
 from typing import Optional
 
 from pydomo import Domo
@@ -24,11 +23,15 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.domoDashboardConnection import (
     DomoDashboardConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: DomoDashboardConnection) -> OMPyDomoClient:
@@ -57,7 +60,8 @@ def test_connection(
     client: OMPyDomoClient,
     service_connection: DomoDashboardConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -72,9 +76,10 @@ def test_connection(
         "GetCharts": client.custom.test_list_cards,
     }
 
-    test_connection_steps(
+    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/dashboard/lightdash/connection.py

@@ -12,7 +12,6 @@
 """
 Source connection handler
 """
-
 from typing import Optional
 
 from metadata.generated.schema.entity.automations.workflow import (
@@ -21,12 +20,16 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.lightdashConnection import (
     LightdashConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.lightdash.client import LightdashApiClient
+from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -49,7 +52,8 @@ def test_connection(
     client: LightdashApiClient,
     service_connection: LightdashConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -60,9 +64,10 @@ def test_connection(
 
     test_fn = {"GetDashboards": custom_executor}
 
-    test_connection_steps(
+    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/dashboard/looker/connection.py

@@ -24,8 +24,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.lookerConnection import (
     LookerConnection,
 )
+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
 
 
 def get_connection(connection: LookerConnection) -> Looker40SDK:
@@ -49,7 +53,8 @@ def test_connection(
     client: Looker40SDK,
     service_connection: LookerConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -76,9 +81,10 @@ def test_connection(
         "ListLookMLModels": list_datamodels_test,
     }
 
-    test_connection_steps(
+    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/dashboard/metabase/connection.py

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.metabaseConnection import (
     MetabaseConnection,
 )
+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.ingestion.source.dashboard.metabase.client import MetabaseClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: MetabaseConnection) -> MetabaseClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: MetabaseClient,
     service_connection: MetabaseConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -49,9 +54,10 @@ def test_connection(
 
     test_fn = {"GetDashboards": custom_executor}
 
-    test_connection_steps(
+    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/dashboard/mode/client.py

@@ -67,15 +67,27 @@ class ModeApiClient:
         )
         self.client = REST(client_config)
 
-    def fetch_all_reports(self, workspace_name: str) -> Optional[list]:
+    def fetch_all_reports(
+        self, workspace_name: str, filter: Optional[str] = "all"
+    ) -> Optional[list]:
         """Method to fetch all reports for Mode
         Args:
             workspace_name:
+            filter:
         Returns:
             dict
         """
+        if filter not in ["custom", "all"]:
+            logger.warning(
+                "Invalid value for filter. Should be one of ['custom', 'all']"
+            )
+            return
+
         all_reports = []
-        response_collections = self.client.get(f"/{workspace_name}/{COLLECTIONS}")
+        filter_param = f"?filter={filter}"
+        response_collections = self.client.get(
+            f"/{workspace_name}/{COLLECTIONS}{filter_param}"
+        )
         collections = response_collections[EMBEDDED]["spaces"]
         for collection in collections:
             response_reports = self.get_all_reports_for_collection(
@@ -175,3 +187,18 @@ class ModeApiClient:
             logger.warning(f"Error fetching all data sources: {exc}")
 
         return None
+
+    def get_workspace(self, workspace_name: str) -> Optional[dict]:
+        """Method to get info about a workspace
+        Args:
+            workspace_name:
+        Returns:
+            dict
+        """
+        try:
+            response = self.client.get(f"/{workspace_name}")
+            return response
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error testing workspace connection: {exc}")
+            raise exc

metadata/ingestion/source/dashboard/mode/connection.py

@@ -21,9 +21,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.modeConnection import (
     ModeConnection,
 )
+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.ingestion.source.dashboard.mode.client import ModeApiClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: ModeConnection) -> ModeApiClient:
@@ -38,7 +42,8 @@ def test_connection(
     client: ModeApiClient,
     service_connection: ModeConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -46,13 +51,14 @@ def test_connection(
 
     test_fn = {
         "CheckDashboards": partial(
-            client.fetch_all_reports, service_connection.workspaceName
+            client.get_workspace, service_connection.workspaceName
         )
     }
 
-    test_connection_steps(
+    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/dashboard/mode/metadata.py

@@ -62,6 +62,7 @@ class ModeSource(DashboardServiceSource):
     ):
         super().__init__(config, metadata)
         self.workspace_name = config.serviceConnection.root.config.workspaceName
+        self.filter_query_param = config.serviceConnection.root.config.filterQueryParam
         self.data_sources = self.client.get_all_data_sources(self.workspace_name)
 
     @classmethod
@@ -80,7 +81,9 @@ class ModeSource(DashboardServiceSource):
         """
         Get List of all dashboards
         """
-        return self.client.fetch_all_reports(self.workspace_name)
+        # If filter param field was empty, we will default to passing "all" to the API
+        filter_param = "all" if not self.filter_query_param else self.filter_query_param
+        return self.client.fetch_all_reports(self.workspace_name, filter_param)
 
     def get_dashboard_name(self, dashboard: dict) -> str:
         """

metadata/ingestion/source/dashboard/mstr/client.py

@@ -30,7 +30,6 @@ from metadata.ingestion.source.dashboard.mstr.models import (
     MstrSearchResult,
     MstrSearchResultList,
 )
-from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -45,8 +44,8 @@ class MSTRClient:
 
     def _get_base_url(self, path=None):
         if not path:
-            return f"{clean_uri(self.config.hostPort)}/{API_VERSION}/"
-        return f"{clean_uri(self.config.hostPort)}/{API_VERSION}/{path}"
+            return f"{self.config.hostPort}/{API_VERSION}/"
+        return f"{self.config.hostPort}/{API_VERSION}/{path}"
 
     def _get_mstr_session(self) -> MSTRRESTSession:
         try:

metadata/ingestion/source/dashboard/mstr/connection.py

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.mstrConnection import (
     MstrConnection,
 )
+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.ingestion.source.dashboard.mstr.client import MSTRClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: MstrConnection) -> MSTRClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: MSTRClient,
     service_connection: MstrConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetProjects": client.get_projects_list}
 
-    test_connection_steps(
+    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/dashboard/powerbi/connection.py

@@ -20,6 +20,9 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.powerBIConnection import (
     PowerBIConnection,
 )
+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.ingestion.source.dashboard.powerbi.client import (
@@ -27,6 +30,7 @@ from metadata.ingestion.source.dashboard.powerbi.client import (
     PowerBiClient,
 )
 from metadata.ingestion.source.dashboard.powerbi.file_client import PowerBiFileClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: PowerBIConnection) -> PowerBiApiClient:
@@ -46,16 +50,18 @@ def test_connection(
     client: PowerBiClient,
     service_connection: PowerBIConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
     """
     test_fn = {"GetDashboards": client.api_client.fetch_dashboards}
 
-    test_connection_steps(
+    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/dashboard/qlikcloud/connection.py

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.qlikCloudConnection import (
     QlikCloudConnection,
 )
+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.ingestion.source.dashboard.qlikcloud.client import QlikCloudClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: QlikCloudConnection) -> QlikCloudClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: QlikCloudClient,
     service_connection: QlikCloudConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetDashboards": client.get_dashboards_list_test_conn}
 
-    test_connection_steps(
+    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/dashboard/qliksense/connection.py

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.qlikSenseConnection import (
     QlikSenseConnection,
 )
+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.ingestion.source.dashboard.qliksense.client import QlikSenseClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: QlikSenseConnection) -> QlikSenseClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: QlikSenseClient,
     service_connection: QlikSenseConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetDashboards": client.get_dashboard_for_test_connection}
 
-    test_connection_steps(
+    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/dashboard/quicksight/connection.py

@@ -22,8 +22,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.dashboard.quickSightConnection import (
     QuickSightConnection,
 )
+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
 
 
 def get_connection(connection: QuickSightConnection):
@@ -43,7 +47,8 @@ def test_connection(
     client: AWSClient,
     service_connection: QuickSightConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> 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
@@ -55,9 +60,10 @@ def test_connection(
         )
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )