openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc3__py3-none-any.whl

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/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/dashboard/dashboard_service.py

@@ -479,6 +479,7 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
         to_entity: Union[Dashboard, DashboardDataModel],
         from_entity: Union[Table, DashboardDataModel, Dashboard],
         column_lineage: List[ColumnLineage] = None,
+        sql: Optional[str] = None,
     ) -> Optional[Either[AddLineageRequest]]:
         if from_entity and to_entity:
             return Either(
@@ -494,6 +495,7 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
                         ),
                         lineageDetails=LineageDetails(
                             source=LineageSource.DashboardLineage,
+                            sqlQuery=sql,
                             columnsLineage=column_lineage,
                         ),
                     )

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

@@ -0,0 +1,268 @@
+#  Copyright 2023 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.
+"""
+REST Auth & Client for MicroStrategy
+"""
+import traceback
+from typing import List, Optional
+
+import requests
+
+from metadata.generated.schema.entity.services.connections.dashboard.microStrategyConnection import (
+    MicroStrategyConnection,
+)
+from metadata.ingestion.connections.test_connections import SourceConnectionException
+from metadata.ingestion.ometa.client import REST, ClientConfig
+from metadata.ingestion.source.dashboard.microstrategy.models import (
+    AuthHeaderCookie,
+    MstrDashboard,
+    MstrDashboardDetails,
+    MstrDashboardList,
+    MstrProject,
+    MstrProjectList,
+    MstrSearchResult,
+    MstrSearchResultList,
+)
+from metadata.utils.helpers import clean_uri
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+API_VERSION = "MicroStrategyLibrary/api"
+HEADERS = {"Content-Type": "application/json"}
+APPLICATION_TYPE = 35
+
+
+class MicroStrategyClient:
+    """
+    Client Handling API communication with MicroStrategy
+    """
+
+    def _get_base_url(self, path=None):
+        if not path:
+            return f"{clean_uri(str(self.config.hostPort))}/{API_VERSION}"
+        return f"{clean_uri(str(self.config.hostPort))}/{API_VERSION}/{path}"
+
+    def __init__(
+        self,
+        config: MicroStrategyConnection,
+    ):
+        self.config = config
+
+        self.auth_params = self._get_auth_header_and_cookies()
+
+        client_config = ClientConfig(
+            base_url=clean_uri(str(self.config.hostPort)),
+            api_version=API_VERSION,
+            extra_headers=self.auth_params.auth_header if self.auth_params else None,
+            allow_redirects=True,
+            cookies=self.auth_params.auth_cookies if self.auth_params else None,
+        )
+
+        self.client = REST(client_config)
+
+    def get_auth_params(self) -> AuthHeaderCookie:
+        """
+        Test whether we can fetch auth_token from the api
+        """
+        data = {
+            "username": self.config.username,
+            "password": self.config.password.get_secret_value(),
+            "loginMode": int(self.config.loginMode),
+            "applicationType": APPLICATION_TYPE,
+        }
+        response = requests.post(
+            url=self._get_base_url("auth/login"), json=data, headers=HEADERS, timeout=60
+        )
+        response.raise_for_status()
+        if (
+            not response.ok
+            or response.status_code != 204
+            or "X-MSTR-AuthToken" not in response.headers
+        ):
+            raise SourceConnectionException(
+                f"Failed to Fetch Token, please validate your credentials and login_mode : {response.text}"
+            )
+        return AuthHeaderCookie(
+            auth_header=response.headers, auth_cookies=response.cookies
+        )
+
+    def _get_auth_header_and_cookies(self) -> Optional[AuthHeaderCookie]:
+        """
+        Send a request to authenticate the user and get headers and
+
+        To know about the data params below please visit
+        https://demo.microstrategy.com/MicroStrategyLibrary/api-docs/index.html#/Authentication/postLogin
+        """
+        try:
+            auth_data = self.get_auth_params()
+            if auth_data:
+                self._set_api_session(auth_data)
+                return auth_data
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Failed to fetch the auth header and cookies due to : [{exc}]"
+            )
+        return None
+
+    def _set_api_session(self, auth_data: AuthHeaderCookie) -> bool:
+        """
+        Set the user api session to active this will keep the connection alive
+        """
+        api_session = requests.put(
+            url=self._get_base_url("sessions"),
+            headers=auth_data.auth_header,
+            cookies=auth_data.auth_cookies,
+            timeout=60,
+        )
+        if api_session.ok:
+            logger.info(
+                f"Connection Successful User {self.config.username} is Authenticated"
+            )
+            return True
+        raise requests.ConnectionError(
+            "Connection Failed, Failed to set an api session, Please validate the credentials"
+        )
+
+    def close_api_session(self) -> None:
+        """
+        Closes the active api session
+        """
+        try:
+            close_api_session = self.client.post(
+                path="/auth/logout",
+            )
+            if close_api_session.ok:
+                logger.info("API Session Closed Successfully")
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to close the api sesison due to [{exc}]")
+
+    def is_project_name(self) -> bool:
+        return bool(self.config.projectName)
+
+    def get_projects_list(self) -> List[MstrProject]:
+        """
+        Get List of all projects
+        """
+        try:
+            resp_projects = self.client.get(
+                path="/projects",
+            )
+
+            project_list = MstrProjectList(projects=resp_projects)
+            return project_list.projects
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to fetch the project list due to [{exc}]")
+
+        return []
+
+    def get_project_by_name(self) -> Optional[MstrProject]:
+        """
+        Get Project By Name
+        """
+        try:
+            resp_projects = self.client.get(
+                path=f"/projects/{self.config.projectName}",
+            )
+
+            project = MstrProject.model_validate(resp_projects)
+            return project
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the project list")
+
+        return None
+
+    def get_search_results_list(
+        self, project_id, object_type
+    ) -> List[MstrSearchResult]:
+        """
+        Get Search Results
+
+        To know about the data params below please visit
+        https://demo.microstrategy.com/MicroStrategyLibrary/api-docs/index.html?#/Browsing/doQuickSearch
+        """
+        try:
+            data = {
+                "project_id": project_id,
+                "type": object_type,
+                "getAncestors": False,
+                "offset": 0,
+                "limit": -1,
+                "certifiedStatus": "ALL",
+                "isCrossCluster": False,
+                "result.hidden": False,
+            }
+            resp_results = self.client.get(
+                path="/searches/results",
+                data=data,
+            )
+
+            results_list = MstrSearchResultList.model_validate(resp_results).result
+            return results_list
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the Search Result list")
+
+        return []
+
+    def get_dashboards_list(self, project_id, project_name) -> List[MstrDashboard]:
+        """
+        Get Dashboard
+        """
+        try:
+            results = self.get_search_results_list(
+                project_id=project_id, object_type=55
+            )
+
+            dashboards = []
+            for result in results:
+                dashboards.append(
+                    MstrDashboard(projectName=project_name, **result.model_dump())
+                )
+
+            dashboards_list = MstrDashboardList(dashboards=dashboards)
+            return dashboards_list.dashboards
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the dashboard list")
+
+        return []
+
+    def get_dashboard_details(
+        self, project_id, project_name, dashboard_id
+    ) -> Optional[MstrDashboardDetails]:
+        """
+        Get Dashboard Details
+        """
+        try:
+            headers = {"X-MSTR-ProjectID": project_id} | self.auth_params.auth_header
+            resp_dashboard = self.client._request(  # pylint: disable=protected-access
+                "GET", path=f"/v2/dossiers/{dashboard_id}/definition", headers=headers
+            )
+
+            return MstrDashboardDetails(
+                projectId=project_id, projectName=project_name, **resp_dashboard
+            )
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to fetch the dashboard with id: {dashboard_id}")
+
+        return None

metadata/ingestion/source/dashboard/{mstr → microstrategy}/connection.py

@@ -17,43 +17,43 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
-from metadata.generated.schema.entity.services.connections.dashboard.mstrConnection import (
-    MstrConnection,
+from metadata.generated.schema.entity.services.connections.dashboard.microStrategyConnection import (
+    MicroStrategyConnection,
 )
 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
+from metadata.ingestion.source.dashboard.microstrategy.client import MicroStrategyClient
 
 
-def get_connection(connection: MstrConnection) -> MSTRClient:
+def get_connection(connection: MicroStrategyConnection) -> MicroStrategyClient:
     """
     Create connection
     """
-    return MSTRClient(connection)
+    return MicroStrategyClient(connection)
 
 
 def test_connection(
     metadata: OpenMetadata,
-    client: MSTRClient,
-    service_connection: MstrConnection,
+    client: MicroStrategyClient,
+    service_connection: MicroStrategyConnection,
     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
     """
 
-    test_fn = {"GetProjects": client.get_projects_list}
+    test_fn = {
+        "GetToken": client.get_auth_params,
+        "GetProjects": client.get_projects_list,
+    }
 
     return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
-        timeout_seconds=timeout_seconds,
     )